#!/usr/bin/env python3
"""
Square lattice of half spheres on substrate.
Test to investigate intensity changes with different numbers of slices.
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm
import matplotlib.pyplot as plt


def get_sample(n_slices=1):
    # Materials (zero absorption)
    particle_color = (0.86, 0.24, 0.18)
    particle_mat = ba.RefractiveMaterial("Particle", particle_color, 3e-05, 0)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-06, 0)
    vacuum = ba.Vacuum()

    # Particles
    ff = ba.SphericalSegment(5*nm, 0, 5*nm)
    particle = ba.Particle(particle_mat, ff)

    # 2D lattices
    lattice = ba.BasicLattice2D(10*nm, 10*nm, 90*deg, 0)

    # Interference functions
    layout = ba.Crystal2D(particle, lattice)
    profile = ba.Profile2DCauchy(100*nm, 100*nm, 0)
    layout.setDecayFunction(profile)

    # Layers
    layer_1 = ba.Layer(vacuum)
    if n_slices > 1:
        layer_1.setNumberOfSlices(n_slices)
    layer_1.deposit2D(layout)
    layer_2 = ba.Layer(substrate_mat)

    # Sample
    sample = ba.Sample()
    sample.addLayer(layer_1)
    sample.addLayer(layer_2)

    return sample


def simulate_horizontal_cut(n_slices=1):
    """Simulate with Nx1 detector for horizontal cut at alpha_f=0.92 deg."""
    sample = get_sample(n_slices)
    beam = ba.Beam(1e9, 0.1*nm, 0.2*deg)
    n = 200
    # Single bin centered at alpha_f=0.92 deg
    detector = ba.SphericalDetector(n, -2*deg, 2*deg, 1, 0.91*deg, 0.93*deg)
    simulation = ba.ScatteringSimulation(beam, sample, detector)
    return simulation.simulate()


def simulate_vertical_cut(n_slices=1):
    """Simulate with 1xN detector for vertical cut at phi_f=0 deg."""
    sample = get_sample(n_slices)
    beam = ba.Beam(1e9, 0.1*nm, 0.2*deg)
    n = 200
    # Single bin centered at phi_f=0 deg
    detector = ba.SphericalDetector(1, -0.01*deg, 0.01*deg, n, 0, 2*deg)
    simulation = ba.ScatteringSimulation(beam, sample, detector)
    return simulation.simulate()


def simulate_depthprobe(n_slices=1):
    """Simulate intensity vs depth."""
    sample = get_sample(n_slices)
    n = 100
    alpha_f = 0.03*deg
    scan = ba.AlphaScan([alpha_f])
    scan.setWavelength(0.1*nm)
    z_axis = ba.EquiDivision("z (nm)", n, -20*nm, 120*nm)
    simulation = ba.DepthprobeSimulation(scan, sample, z_axis)
    return simulation.simulate(), alpha_f/deg


# Color cycle for different N values
colors = ['C0', 'C1', 'C2', 'C3']


if __name__ == '__main__':
    # Slice counts to compare
    slice_counts = [1, 2, 5, 20]

    # Create figure with 1 row, 3 columns
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    # Horizontal cuts at alpha_f=0.92 deg
    ax = axes[0]
    for i, n in enumerate(slice_counts):
        result = simulate_horizontal_cut(n_slices=n)
        ax.semilogy(result.xCenters(), result.flatVector(),
                    label=f'N={n}', color=colors[i])
    ax.set_title('Horizontal cut at alpha_f=0.92 deg')
    ax.set_xlabel('phi_f (deg)')
    ax.set_ylabel('Intensity')
    ax.legend()

    # Vertical cuts at phi_f=0 deg
    ax = axes[1]
    for i, n in enumerate(slice_counts):
        result = simulate_vertical_cut(n_slices=n)
        ax.semilogy(result.yCenters(), result.flatVector(),
                    label=f'N={n}', color=colors[i])
    ax.set_title('Vertical cut at phi_f=0 deg')
    ax.set_xlabel('alpha_f (deg)')
    ax.set_ylabel('Intensity')
    ax.legend()

    # Intensity vs depth
    ax = axes[2]
    for i, n in enumerate(slice_counts):
        result, alpha_f = simulate_depthprobe(n_slices=n)
        ax.semilogy(result.yCenters(), result.flatVector(),
                    label=f'N={n}', color=colors[i])
    ax.set_title(f'Intensity vs depth at alpha_f={alpha_f:.4f} deg')
    ax.set_xlabel('z (nm)')
    ax.set_ylabel('Intensity')
    ax.legend()

    plt.tight_layout()
    plotargs = bp.parse_commandline()
    bp.export(**plotargs)
