#!/usr/bin/env python3
"""
Intensity vs angle at fixed depth z=100nm (upper vacuum boundary).
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm

#  beam data
ai_min = 0*deg
ai_max = 1*deg
wl = 0.03*nm  # wavelength

#  fixed depth at upper boundary
z_fixed = 100*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 scanning over angles at fixed depth.
    """
    nz = 1  # Single depth point
    na = 500

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

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

    return simulation


if __name__ == '__main__':
    sample = get_sample()
    simulation = get_simulation(sample)
    result = simulation.simulate()
    plotargs = bp.parse_commandline()
    import numpy as np

    # Extract 1D data
    data = result.intensities()
    intensity = data.flatten()
    alpha_axis = result.axis(0)
    alpha_values = np.linspace(alpha_axis.min(), alpha_axis.max(),
                               alpha_axis.size()) / deg

    # Custom plot for intensity vs angle
    bp.plt.figure(figsize=(10, 6))
    ax = bp.plt.gca()
    ax.plot(alpha_values, intensity, 'b-', linewidth=2,
            label=f'z = {z_fixed/nm:.0f} nm')
    ax.axhline(y=1, color='r', linestyle='--', linewidth=1,
               label='Expected vacuum intensity')

    # Calculate and mark critical angle
    delta = 1e-5
    theta_c = np.sqrt(2*delta) * 180/np.pi  # in degrees
    ax.axvline(x=theta_c, color='orange', linestyle=':', linewidth=1.5,
               label=f'Critical angle θc = {theta_c:.3f}°')

    ax.set_xlabel('Incident angle αi (°)', fontsize=14)
    ax.set_ylabel('Intensity', fontsize=14)
    title = f'Intensity vs Angle at z = {z_fixed/nm:.0f} nm (upper boundary)'
    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.export(**plotargs)
