import bornagain as ba
from bornagain import ba_plot as bp, deg, nm
import numpy as np

def get_sample(P):
    """
    A sample with cylinders and pyramids on a substrate,
    forming a hexagonal lattice.
    """
    radius = P['radius']
    lattice_length = P['length']

    vacuum = ba.Vacuum()
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-6, 2e-8)
    particle_color = (0.86, 0.24, 0.18)
    particle_mat = ba.RefractiveMaterial("Particle", particle_color, 6e-4, 2e-8)

    sphere_ff = ba.Sphere(radius)
    sphere = ba.Particle(particle_mat, sphere_ff)

    layout = ba.Crystal2D(sphere, ba.HexagonalLattice2D(lattice_length, 0))
    profile = ba.Profile2DCauchy(100*nm, 100*nm, 0)
    layout.setDecayFunction(profile)

    vacuum_layer = ba.Layer(vacuum)
    vacuum_layer.deposit2D(layout)
    substrate_layer = ba.Layer(substrate_mat)

    sample = ba.Sample()
    sample.addLayer(vacuum_layer)
    sample.addLayer(substrate_layer)
    return sample


def get_simulation(P):
    """
    Create and return GISAXS simulation with beam and detector defined
    """
    n = 100
    beam = ba.Beam(1e8, 0.1*nm, 0.2*deg)
    detector = ba.SphericalDetector(n, -1*deg, 1*deg, n, 0, 2*deg)
    sample = get_sample(P)
    simulation = ba.ScatteringSimulation(beam, sample, detector)
    simulation.options().setUseAvgMaterials(False)
    return simulation


def fake_data():
    """
    Generating "real" data by adding noise to the simulated data.
    """
    P = {'radius': 6*nm, 'length': 12*nm}
    simulation = get_simulation(P)
    result = simulation.simulate()

    return result.noisy(0.1, 0.1)

if __name__ == '__main__':
    raise Exception("This source file is a module for use from fitting examples."
                    " It is not meant to be called as a main program.")
