#!/usr/bin/env python3
"""
Fake GISAS data for testing and benchmarking 2D fits.
The sample model basically consists of dilute cylinders on a substrate.
There are, however, some systematic distortions:
  * Refractive indices changed, and absorption strongly enhanced;
  * Cylinders replaced by a mixture of cones and segmented spheroids;
  * Beam wavelength and alpha_i slightly changed.
  * Detector x-axis skewed;
"""

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

mat_vacuum = ba.RefractiveMaterial("Vacuum", 0, 0)
# nominal material constants are (6e-6, 2e-8) and (6e-4, 2e-8)
mat_substrate = ba.RefractiveMaterial("Substrate", 6.3e-6, 4e-7)
mat_particle = ba.RefractiveMaterial("Particle", 3e-5, 2e-7)


def get_sample(params):
    h = params["cylinder_height"]
    r = params["cylinder_radius"]
    ro = params["surface_density"]

    ff1 = ba.Cone(1.1*r, h, 80*deg)
    ff2 = ba.TruncatedSpheroid(r, 3.3*h, 2.5*h/r, 2.2*h)

    layout = ba.ParticleLayout()
    layout.addParticle(ba.Particle(mat_particle, ff1), .4)
    layout.addParticle(ba.Particle(mat_particle, ff2), .6)
    layout.setTotalParticleSurfaceDensity(ro)

    layer_1 = ba.Layer(mat_vacuum)
    layer_1.addLayout(layout)
    layer_2 = ba.Layer(mat_substrate)

    sigma, hurst, corrLength = 0.2*nm, 0.3, 20*nm
    roughness = ba.LayerRoughness(sigma, hurst, corrLength)

    sample = ba.Sample()
    sample.addLayer(layer_1)
    sample.addLayerWithTopRoughness(layer_2, roughness)
    return sample


def get_simulation(params):
    # nominal beam parameters are 0.1*nm, 0.2*deg
    beam = ba.Beam(10**params['lg(intensity)'], 0.0993*nm, 0.202*deg)

    # nominal detector x-axis starts from -1.5*deg
    det = ba.SphericalDetector(120, -1.49*deg, 1.5*deg, 120, 0, 3*deg)

    sample = get_sample(params)

    simulation = ba.ScatteringSimulation(beam, sample, det)
    if 'lg(background)' in params:
        simulation.setBackground(
            ba.ConstantBackground(10**params['lg(background)']))

    return simulation


def model_parameters():
    return {
        'lg(intensity)': 9,
        'lg(background)': 1.2,
        'cylinder_height': 4*nm,
        'cylinder_radius': 5*nm,
        'surface_density': 0.005
    }


def fake_data():
    """
    Generate fake "experimental" data, and save them as numpy array.
    """

    params = model_parameters()

    # Compute model distribution
    simulation = get_simulation(params)
    result = simulation.simulate()
    bp.plot_simulation_result(result)

    theory = result.npArray()

    # Draw noisy data
    data = np.random.poisson(theory)

    # Save to numpy
    np.savetxt("faked-gisas.txt", data)


if __name__ == '__main__':
    np.random.seed(1)
    fake_data()
