#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function

# We'll need numpy and matplotlib for plotting our results
import numpy as np
import matplotlib.pyplot as plt

# A least squares fitting algorithm from scipy
from scipy.optimize.minpack import leastsq

# DiffPy-CMI modules for building a fitting recipe
from diffpy.Structure import loadStructure
from diffpy.srfit.pdf import PDFContribution
from diffpy.srfit.fitbase import FitRecipe, FitResults

# Files containing our experimental data and structure file
dataFile = "NOM_9999_Si_d_new_NewCapA_ftfrgr.gr"
structureFile = "si.cif"

# The first thing to construct is a contribution. Since this is a simple
# example, the contribution will simply contain our PDF data and an associated
# structure file. We'll give it the name "silicon"
siPDF = PDFContribution("silicon")
siPDF.setScatteringType(type = 'N')

# Load the data and set the r-range over which we'll fit
siPDF.loadData(dataFile)
siPDF.setCalculationRange(xmin = 1, xmax = 50.0, dx = 0.01)

# Add the structure from our cif file to the contribution
siStructure = loadStructure(structureFile)
siPDF.addStructure("silicon", siStructure)


# The FitRecipe does the work of calculating the PDF with the fit variable
# that we give it.
siFit = FitRecipe()

# give the PDFContribution to the FitRecipe
siFit.addContribution(siPDF)

# Configure the fit variables and give them to the recipe.
lat = siPDF.silicon.phase.getLattice()
siFit.newVar("a", lat.a.getValue())
siFit.constrain(lat.a, "a")
siFit.constrain(lat.b, "a")
siFit.constrain(lat.c, "a")

atoms = siPDF.silicon.phase.getScatterers()
siFit.newVar("Si_Uiso", 0.006)
for atom in atoms:
	if atom.element.title() == "Si":
		siFit.constrain(atom.Uiso, "Si_Uiso")


# As usual, we add variables for the overall scale of the PDF and a delta2
# parameter for correlated motion of neighboring atoms.
siFit.addVar(siPDF.scale, 1)
siFit.addVar(siPDF.silicon.delta2, 2)

# We refine Qdamp and Qbroad based on prior information about our beamline.
siFit.addVar(siPDF.qdamp, 0.0176590)
siFit.addVar(siPDF.qbroad, 0.0191822)

# Turn off printout of iteration number.
siFit.clearFitHooks()

# We can now execute the fit using scipy's least square optimizer.
print("Refine PDF using scipy's least-squares optimizer:")
print("  variables:", siFit.names)
print("  initial values:", siFit.values)
leastsq(siFit.residual, siFit.values)
print("  final values:", siFit.values)
print()

# Obtain and display the fit results.
siResults = FitResults(siFit)
print("FIT RESULTS\n")
print(siResults)

# Save the fitted structure
siStructure.write("si_fitted.cif", format = "cif")

# Save the fitted data
siFit.silicon.profile.savetxt("si_fitted.fgr")


# Plot the observed and refined PDF.
# Get the experimental data from the recipe
r = siFit.silicon.profile.x
gobs = siFit.silicon.profile.y

# Get the calculated PDF and compute the difference between the calculated and
# measured PDF
gcalc = siFit.silicon.evaluate()
baseline = 1.1 * gobs.min()
gdiff = gobs - gcalc

# Plot!
plt.figure()
plt.plot(r, gobs, 'bo', label="G(r) data",
         markerfacecolor='none', markeredgecolor='b')
plt.plot(r, gcalc, 'r-', label="G(r) fit")
plt.plot(r, gdiff + baseline, 'g-', label="G(r) diff")
plt.plot(r, np.zeros_like(r) + baseline, 'k:')
plt.xlabel(r"r ($\AA$)")
plt.ylabel(r"G ($\AA^{-2}$)")
plt.legend()

plt.show()
