#!/usr/bin/env python3
"""
Generates random rough surface of the sample with the given
heights statistics and lateral roughness spectrum.
"""

import numpy as np
import bornagain as ba
from bornagain import nm, nm2, nm3
from matplotlib import gridspec
from bornagain.numpyutil import Arrayf64Converter as dac

def plot(h):
    rms = np.std(surface)

    fig = bp.plt.figure(figsize=(10, 6))
    gs = gridspec.GridSpec(3, 2)
    gs.update(left=0.1, right=0.48, wspace=0.15, hspace=0.5)

    # map
    fig.add_subplot(gs[:-1, :])
    bp.plt.imshow(h, extent = [0, Lx, 0, Ly], cmap='inferno', aspect='auto')
    bp.plt.colorbar()

    # psd
    fig.add_subplot(gs[-1, 0])
    factor = X_points*Y_points/np.sqrt(Lx*Ly)
    N = np.floor(X_points/2 + 1).astype(int)
    psd = abs(np.fft.rfft2(h)/factor)**2
    psd_points_x = np.array([psd[0][i] for i in range(1, N)])
    f_points_x = np.array([i/Lx for i in range(1, N)])
    bp.plt.plot(f_points_x, psd_points_x)
    bp.plt.yscale('log')
    bp.plt.xscale('log')
    bp.plt.title('PSD')
    bp.plt.xlabel('$\\nu, nm^{-1}$')

    # hist
    fig.add_subplot(gs[-1, 1])
    bp.plt.hist(h.flatten(), range=[-4*rms, 4*rms],
             bins=300, color='black')
    bp.plt.title('Height histogram')
    bp.plt.xlabel('h, nm')
    bp.plt.tick_params(left = False, labelleft = False)

    bp.plt.show()

def get_sample():
    """
    A sample with one layer on a substrate, with partially
    correlated roughnesses.
    """
    # defining materials
    vacuum = ba.RefractiveMaterial("ambience", 0, 0)
    material_layer = ba.RefractiveMaterial("layer", 5e-6, 0)
    material_substrate = ba.RefractiveMaterial("substrate", 15e-6, 0)

    # defining roughness
    height_distribution = ba.ErfInterlayer()
    max_spat_freq = 0.5*nm

    # for layer
    omega = 1*nm3
    a1 = 1
    a2 = 10*nm
    a3 = 100*nm2
    a4 = 1000*nm3
    autocorr_layer = ba.LinearGrowthModel(omega, a1, a2, a3, a4, max_spat_freq)
    roughness_layer = ba.LayerRoughness(autocorr_layer, height_distribution)

    # for substrate
    sigma = 0.9*nm
    alpha = 0.5
    xi = 35*nm
    autocorr_sub = ba.K_CorrelationModel(sigma, alpha, xi, max_spat_freq)
    roughness_sub = ba.LayerRoughness(autocorr_sub, height_distribution)

    # defining layers
    l_ambience = ba.Layer(vacuum)
    l_layer = ba.Layer(material_layer, 25*nm, roughness_layer)
    l_substrate = ba.Layer(material_substrate, roughness_sub)

    # adding layers
    my_sample = ba.Sample()
    my_sample.addLayer(l_ambience)
    my_sample.addLayer(l_layer)
    my_sample.addLayer(l_substrate)

    return my_sample

if __name__ == '__main__':

    # sample size
    Lx = 1000*nm
    Ly = 1000*nm
    X_points = 10
    Y_points = 10

    sample = get_sample()
    layer_index = 1 # top interface
    seed = -1 # seed < 0 ==> random every time

    # generate roughness map
    roughness_map = ba.RoughnessMap(X_points, Y_points, Lx, Ly,
                                    sample, layer_index, seed)
    surface = dac.npArray(roughness_map.generate())
    print("rms = {:.3}".format(np.std(surface)),"nm")

    from bornagain import ba_plot as bp
    plotargs = bp.parse_commandline()
    bp.export(**plotargs)
