#!/usr/bin/env python3
"""
Fitting experimental data: spherical nanoparticles with size distribution
in 3 layers system (experiment at GALAXI).
"""
import os
import bornagain as ba
from bornagain import deg, nm, ba_fitmonitor, R3
from matplotlib import pyplot as plt

wavelength = 1.34*ba.angstrom
alpha_i = 0.463*ba.deg

# detector setup as given from instrument responsible
pilatus_npx, pilatus_npy = 981, 1043
# pilatus_pixel_size = 0.172  # in mm
# detector_distance = 1730.0  # in mm
# beam_xpos, beam_ypos = 597.1, 323.4  # in pixels

datadir = os.getenv('BA_DATA_DIR', '')
if not datadir:
    raise Exception("Environment variable BA_DATA_DIR not set")

# sample model

radius = 5.75*ba.nm
sigma = 0.4
distance = 53.6*ba.nm
disorder = 10.5*ba.nm
kappa = 17.5
ptfe_thickness = 22.1*ba.nm
hmdso_thickness = 18.5*ba.nm

def get_sample(P):
    radius = P["radius"]
    sigma = P["sigma"]
    distance = P["distance"]

    # defining materials
    vacuum = ba.RefractiveMaterial("Vacuum", 0, 0)
    material_Si = ba.RefractiveMaterial("Si", 5.7816e-6, 1.0229e-7)
    material_Ag = ba.RefractiveMaterial("Ag", 2.2475e-5, 1.6152e-6)
    material_PTFE = ba.RefractiveMaterial("PTFE", 5.20509e-6, 1.9694e-8)
    material_HMDSO = ba.RefractiveMaterial("HMDSO", 2.0888e-6, 1.3261e-8)

    # collection of particles with size distribution
    nparticles = 20
    nfwhm = 2.0
    sphere_ff = ba.Sphere(radius)

    sphere = ba.Particle(material_Ag, sphere_ff)
    position = R3(0, 0, -1*hmdso_thickness)
    sphere.translate(position)
#     ln_distr = ba.DistributionLogNormal(radius, sigma)
#     par_distr = ba.ParameterDistribution(
#         "/Particle/Sphere/Radius", ln_distr, nparticles, nfwhm,
#     part_coll = ba.ParticleDistribution(sphere, par_distr)

    # interference function
    interference = ba.InterferenceRadialParacrystal(
        distance, 1e6*ba.nm)
    interference.setKappa(kappa)
    interference.setDomainSize(2e4*nm)
    pdf = ba.Profile1DGauss(disorder)
    interference.setProbabilityDistribution(pdf)

    # assembling particle layout
    layout = ba.ParticleLayout()
    layout.addParticle(sphere, 1)
    layout.setInterference(interference)
    layout.setTotalParticleSurfaceDensity(1)

    # roughness
    r_ptfe = ba.LayerRoughness(2.3*ba.nm, 0.3, 5*ba.nm)
    r_hmdso = ba.LayerRoughness(1.1*ba.nm, 0.3, 5*ba.nm)

    # layers
    vacuum_layer = ba.Layer(vacuum)
    hmdso_layer = ba.Layer(material_HMDSO, hmdso_thickness)
    hmdso_layer.addLayout(layout)
    ptfe_layer = ba.Layer(material_PTFE, ptfe_thickness)
    substrate_layer = ba.Layer(material_Si)

    # assembling sample
    sample = ba.MultiLayer()
    sample.addLayer(vacuum_layer)
    sample.addLayerWithTopRoughness(hmdso_layer, r_hmdso)
    sample.addLayerWithTopRoughness(ptfe_layer, r_ptfe)
    sample.addLayer(substrate_layer)

    return sample

def create_detector(beam):
    """
    A model of the GALAXY detector
    """
    nx = pilatus_npx
    ny = pilatus_npy
    return ba.Detector2D(3.4*deg, 1.84*deg, nx, ny, 0, 0.32*deg)


def create_simulation(P):
    """
    Creates and returns GISAS simulation with beam and detector defined
    """
    beam = ba.Beam(1.2e7, wavelength, alpha_i)
    sample = get_sample(P)
    detector = create_detector(beam)
    simulation = ba.ScatteringSimulation(beam, sample, detector)

    # simulation.detector().setRegionOfInterest(85, 70, 120, 92.)
    # beamstop:
    # simulation.detector().addMask(ba.Rectangle(101.9, 82.1, 103.7, 85.2), True)

    return simulation


def load_data(filename):
    """
    Loads experimental data and returns numpy array.
    """
    filepath = os.path.join(datadir, filename)
    return ba.readData2D(filepath)


def run_fitting():
    data = load_data("scatter2d/galaxi_data.tif.gz")

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(create_simulation, data, 1)
    fit_objective.initPrint(10)
    observer = ba_fitmonitor.PlotterGISAS()
    plt.close() # (hide plot) fit_objective.initPlot(10, observer)

    P = ba.Parameters()
    P.add("radius", 5.*nm, min=4, max=6, step=0.1*nm)
    P.add("sigma", 0.55, min=0.2, max=0.8, step=0.01)
    P.add("distance", 27.*nm, min=20, max=70)

    minimizer = ba.Minimizer()
    minimizer.setMinimizer("Minuit2", "Migrad", "MaxFunctionCalls=2")
    result = minimizer.minimize(fit_objective.evaluate, P)
    fit_objective.finalize(result)


if __name__ == '__main__':
    run_fitting()
    plt.show()
