#!/usr/bin/env python3
"""
Cylindrical mesocrystal on a substrate with slicing.
Test to investigate intensity changes with different numbers of slices.
"""
import numpy as np
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm, nm2, R3
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, 6e-05, 0)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-06, 0)
    vacuum = ba.Vacuum()

    # Basis particle
    inner_ff = ba.Sphere(3*nm)
    inner_particle = ba.Particle(particle_mat, inner_ff)

    # 3D lattice
    lattice = ba.Lattice3D(R3(8*nm, 0, 0), R3(0, 8*nm, 0),
                           R3(0, 0, 8*nm))

    # Crystal
    crystal = ba.Crystal(inner_particle, lattice)

    # Mesocrystal: shaped crystal
    outer_ff = ba.Cylinder(20*nm, 50*nm)
    outer_particle = ba.Mesocrystal(crystal, outer_ff)

    # Layers
    layer_1 = ba.Layer(vacuum)
    layer_1.deposit2D(ba.Dilute2D(0.0001/nm2, outer_particle))
    if n_slices > 1:
        layer_1.setNumberOfSlices(n_slices)
    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, reciprocal=False):
    """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 = <%= test_mode ? 11 : 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)
    if reciprocal:
        simulation.options().setMesoReciprocalSum(True, 2.5)
    return simulation.simulate()


def simulate_vertical_cut(n_slices=1, reciprocal=False):
    """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 = <%= test_mode ? 11 : 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)
    if reciprocal:
        simulation.options().setMesoReciprocalSum(True, 2.5)
    return simulation.simulate()


def simulate_depthprobe(n_slices=1, reciprocal=False):
    """Simulate intensity vs depth."""
    sample = get_sample(n_slices)
    n = <%= test_mode ? 11 : 100 %>
    alpha_f = 0.001*deg if reciprocal else 0.0064*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)
    if reciprocal:
        simulation.options().setMesoReciprocalSum(True, 2.5)
    return simulation.simulate(), alpha_f/deg


def get_sld_profile(n_slices=1, as_datafield=False):
    """Get SLD profile from resampled sample."""
    sample = get_sample(n_slices)
    n = <%= test_mode ? 11 : 200 %>
    z_min, z_max = -20*nm, 60*nm
    z = np.linspace(z_min, z_max, n)
    sld = ba.materialProfileSLD(sample, n, z_min, z_max)
    sld_real = np.array([s.real for s in sld])
    if as_datafield:
        # Create Datafield for persistence testing
        z_axis = ba.EquiDivision("z (nm)", n, z_min/nm, z_max/nm)
        frame = ba.Frame(z_axis)
        return ba.Datafield(frame, list(sld_real))
    return z/nm, sld_real


# 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 2 rows, 4 columns
    fig, axes = plt.subplots(2, 4, figsize=(18, 10))

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

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

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

    ax = axes[0, 3]
    for i, n in enumerate(slice_counts):
        z, sld = get_sld_profile(n_slices=n)
        ax.plot(z, sld, label=f'N={n}', color=colors[i])
    ax.set_title('Real-space: SLD profile')
    ax.set_xlabel('z (nm)')
    ax.set_ylabel('SLD (real part)')
    ax.legend()

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

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

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

    ax = axes[1, 3]
    for i, n in enumerate(slice_counts):
        z, sld = get_sld_profile(n_slices=n)
        ax.plot(z, sld, label=f'N={n}', color=colors[i])
    ax.set_title('Fourier: SLD profile')
    ax.set_xlabel('z (nm)')
    ax.set_ylabel('SLD (real part)')
    ax.legend()

    plt.tight_layout()
<%- if test_mode -%>
    # Persistence test for N=20 slices
    from bornagain import ba_check
    results = [
        simulate_horizontal_cut(n_slices=20, reciprocal=False),
        simulate_vertical_cut(n_slices=20, reciprocal=False),
        simulate_depthprobe(n_slices=20, reciprocal=False)[0],
        get_sld_profile(n_slices=20, as_datafield=True),
    ]
    ba_check.persistence_test(results)
<%- elsif figure_mode -%>
    plotargs = bp.parse_commandline()
    bp.export(**plotargs)
<%- else -%>
    ba.showSample3D(get_sample(), sample_size=250*nm, seed=0)
    plt.show()
<%- end -%>
