#!/usr/bin/env python3
"""
GISAS from a hexagonal lattice of magnetic domains with random orientation.

The layer is represented by flat hexagonal domain particles whose refractive
contrast equals the matrix contrast. Their magnetization vectors point into
several in-plane directions; the Mixture makes the domain orientation random
from site to site.
"""
import colorsys
import math

import bornagain as ba
from bornagain import ba_plot as bp, deg, nm, R3


def domain_field(angle):
    induction = 1e6
    return R3(induction*math.cos(angle), induction*math.sin(angle), 0)


def get_sample():
    # Materials
    vacuum = ba.Vacuum()
    # Matrix and domains share the same nuclear SLD, so lateral domain
    # scattering is essentially magnetic-only in the selected spin-flip channel.
    domain_delta = 2e-5
    absorption = 1e-9
    matrix_mat = ba.RefractiveMaterial(
        "Matrix", (0.73, 0.76, 0.78), domain_delta, absorption)
    substrate_mat = ba.RefractiveMaterial(
        "Substrate", (0.28, 0.57, 0.82), 6e-6, absorption)

    # Magnetic domain variants
    # Flat-to-flat width of each regular hexagonal domain.
    domain_flat_width = 50*nm
    domain_height = 20*nm
    domain_edge = domain_flat_width/math.sqrt(3)
    domain_ff = ba.Prism6(domain_edge, domain_height)
    n_orientations = 6

    mix = ba.Mixture()
    for i in range(n_orientations):
        angle = i*360*deg/n_orientations
        # hue encodes the magnetization direction in the 3D view
        domain_color = colorsys.hsv_to_rgb(angle/(360*deg), 0.65, 0.72)
        domain_mat = ba.RefractiveMaterial(
            f"Domain{i}", domain_color, domain_delta, absorption,
            domain_field(angle))
        domain = ba.Particle(domain_mat, domain_ff)
        domain.rotate(ba.RotationZ(30*deg))  # facet-to-facet packing
        mix.addParticle(domain, 1)

    # Layers
    fill_fraction = 0.98
    lattice_period = domain_flat_width/math.sqrt(fill_fraction)
    layout = ba.Crystal2D(mix, ba.HexagonalLattice2D(lattice_period, 0*deg))
    layout.setDecayFunction(ba.Profile2DCauchy(2000*nm, 2000*nm, 0))

    layer = ba.Layer(matrix_mat, domain_height)
    layer.deposit2D(layout)

    sample = ba.Sample()
    sample.addLayer(ba.Layer(vacuum))
    sample.addLayer(layer)
    sample.addLayer(ba.Layer(substrate_mat))
    return sample


def get_simulation(sample):
    beam = ba.Beam(1e9, 0.2*nm, 0.5*deg)
    n = 21
    detector = ba.SphericalDetector(n, -2*deg, 2*deg, n, 0, 4*deg)

    beam.setPolarization(R3(0, 0, 1))
    detector.setAnalyzer(R3(0, 0, -1))

    simulation = ba.ScatteringSimulation(beam, sample, detector)
    return simulation


if __name__ == '__main__':
    sample = get_sample()
    simulation = get_simulation(sample)
    result = simulation.simulate()
