#!/usr/bin/env python3
"""
1D depth profile at fixed angle showing standing wave pattern.
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm

#  beam data
ai_fixed = 0.0044*deg  # fixed angle where maximum occurs
wl = 0.03*nm  # wavelength

#  depth position span
z_min = -350*nm
z_max = 250*nm


def get_sample():
    # Define materials
    vac_color = (0.90, 0.93, 0.97)
    vac_mat = ba.RefractiveMaterial("Vac", vac_color, 0, 0)
    a_color = (0.05, 0.62, 0.55)
    a_mat = ba.RefractiveMaterial("A", a_color, 1e-5, 3e-6)

    # Define layers
    layer_top = ba.Layer(vac_mat)
    layer_1 = ba.Layer(a_mat, 300*nm)
    layer_bot = ba.Layer(vac_mat)

    # Define sample
    sample = ba.Sample()
    sample.addLayer(layer_top)
    sample.addLayer(layer_1)
    sample.addLayer(layer_bot)

    return sample


def get_simulation(sample):
    """
    Returns a depth-probe simulation at a single angle.
    """
    nz = 500
    na = 1  # Just one angle

    scan = ba.AlphaScan(na, ai_fixed, ai_fixed)
    scan.setWavelength(wl)
    footprint = ba.FootprintSquare(0.01)
    scan.setFootprint(footprint)

    z_axis = ba.EquiDivision("z", nz, z_min, z_max)
    simulation = ba.DepthprobeSimulation(scan, sample, z_axis)

    return simulation


if __name__ == '__main__':
    sample = get_sample()
    simulation = get_simulation(sample)
    result = simulation.simulate()
    ba.showSample3D(sample, sample_size=500*nm, seed=0)
    import numpy as np

    # Extract 1D data
    data = result.intensities()
    intensity = data.flatten()
    z_axis = result.axis(1)
    z_values = np.linspace(z_axis.min(), z_axis.max(), z_axis.size()) / nm

    # Custom plot for 1D profile
    bp.plt.figure(figsize=(10, 6))
    ax = bp.plt.gca()
    ax.plot(z_values, intensity, 'b-', linewidth=2)
    ax.axvline(x=0, color='gray', linestyle=':', linewidth=1)
    ax.axvline(x=-300, color='gray', linestyle=':', linewidth=1)

    # Shade regions
    ax.axvspan(z_min/nm, -300, alpha=0.1, color='blue', label='Bottom vacuum')
    ax.axvspan(-300, 0, alpha=0.1, color='orange', label='Material A (300nm)')
    ax.axvspan(0, z_max/nm, alpha=0.1, color='blue', label='Top vacuum')

    ax.set_xlabel('Depth z (nm)', fontsize=14)
    ax.set_ylabel('Intensity', fontsize=14)
    title = f'Depth Profile at α = {ai_fixed/deg:.6f}°'
    ax.set_title(title, fontsize=16)
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=10)
    ax.set_ylim([0, max(intensity)*1.1])

    bp.plt.tight_layout()
    bp.plt.show()
