# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import unicode_literals

import unittest
import numpy as np

from pymatgen.util.testing import PymatgenTest
from pymatgen.analysis.defects.core import DefectEntry, Vacancy
from pymatgen.analysis.defects.corrections import FreysoldtCorrection

class DefectsCorrectionsTest(PymatgenTest):
    def test_freysoldt(self):
        #a trivial test for freysoldt correction
        struc = PymatgenTest.get_structure("VO2")
        abc = struc.lattice.abc
        axisdata = [np.arange(0., lattval, 0.01) for lattval in abc]
        bldata = [np.array([1. for u in np.arange(0., lattval, 0.01)])  for lattval in abc]
        dldata = [np.array([(-1-np.cos(2*np.pi*u/lattval)) for u in np.arange(0., lattval, 0.01)])  for lattval in abc]
        params = {'axis_grid':axisdata, 'bulk_planar_averages': bldata, 'defect_planar_averages':dldata}
        vac = Vacancy(struc, struc.sites[0], charge=-3)
        fc = FreysoldtCorrection(15)

        #test electrostatic correction
        es_corr = fc.perform_es_corr(struc.lattice, 15., -3)
        self.assertEqual(es_corr, 2.929238)

        #test potential alignment method
        pot_corr = fc.perform_pot_corr( axisdata[0], bldata[0], dldata[0], struc.lattice, 15., -3, vac.site.coords, 0)
        self.assertEqual(pot_corr, 2.917553311042547)

        #test entry full correction method
        de = DefectEntry(vac, 0., corrections={}, parameters=params, entry_id=None)
        val = fc.get_correction(de)
        self.assertEqual(val['freysoldt_electrostatic'], 2.929238)
        self.assertEqual(val['freysoldt_potential_alignment'], 4.579907163820233)


if __name__ == "__main__":
    unittest.main()

@@ -362,16 +362,17 @@
 ection%0A%0A
+%0A
 class De
@@ -536,16 +536,48 @@
 (%22VO2%22)%0A
+        struc.make_supercell(3)%0A
         
@@ -645,18 +645,17 @@
 tval, 0.
-01
+2
 ) for la
@@ -725,34 +725,33 @@
 (0., lattval, 0.
-01
+2
 )%5D)  for lattval
@@ -856,10 +856,9 @@
 , 0.
-01
+2
 )%5D) 
@@ -1211,24 +1211,24 @@
 s_corr, 
-2.929238
+0.976412
 )%0A%0A     
@@ -1423,23 +1423,23 @@
 , 2.
-917553311042547
+836369987722345
 )%0A%0A 
@@ -1661,16 +1661,16 @@
 '%5D, 
-2.929238
+0.976412
 )%0A  
@@ -1750,16 +1750,1197 @@
 20233)%0A%0A
+        #test the freysoldt plotter%0A        pltsaver = %5B%5D%0A        for ax in range(3):%0A            x = fc.metadata%5B'pot_plot_data'%5D%5Bax%5D%5B'x'%5D%0A            Vr = fc.metadata%5B'pot_plot_data'%5D%5Bax%5D%5B'Vr'%5D%0A            dft_diff = fc.metadata%5B'pot_plot_data'%5D%5Bax%5D%5B'dft_diff'%5D%0A            final_shift = fc.metadata%5B'pot_plot_data'%5D%5Bax%5D%5B'final_shift'%5D%0A            check = fc.metadata%5B'pot_plot_data'%5D%5Bax%5D%5B'check'%5D%0A            fp = freysoldt_plotter( x, Vr, dft_diff, final_shift, check, title='axis_'+str(ax+1), saved=False)%0A            if fp: #if plot exists then append it...%0A                pltsaver.append(fp)%0A        self.assertEqual(len(pltsaver), 3)%0A%0A    def test_kumagai(self):%0A        #a trivial test for kumagai correction%0A%0A        #test find_optimal_gamma method%0A        # gam = find_optimal_gamma(structure, epsilon)%0A%0A        #test generate_g_sum method%0A        # g_sum = generate_g_sum( structure, epsilon, dim, gamma)%0A%0A        # KumagaiCorrection( dielectric_tensor)%0A%0A        #test electrostatic correction%0A%0A        #test potential alignment method%0A%0A        #test entry full correction method (including%0A%0A        #test the kumagai plotter%0A        # kumagai_plotter%0A%0A        pass%0A%0A%0A
 %0Aif __na

update freysoldt unit test for supercell