import numpy as np
import odeg as od
import refs.hartreeref as hr
import os
from time import time
import matplotlib.pyplot as plt

plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300

if __name__ == '__main__':
    
    pos_charge_grid = np.arange(0.25, 22, .25)
    n_spin = 2
    
    cfg = {
        'r0': 0.5,
        'rs': 2,
        'p_max': 5,
        'pos_charge': 10,
        'theta': .125,
        'return_mu': True,
        }
    
    filename = f"f_with_r0_{cfg['r0']}_rs_{cfg['rs']}_pmax_{cfg['p_max']}_theta_{cfg['theta']}.txt"
    res_path = './results/fxc/'
    if not(os.path.exists('./results/')):
        os.mkdir('./results/')
    if not(os.path.exists(res_path)):
        os.mkdir(res_path)
    file = open(res_path + filename, "w")
    
    t_full = 0
    t_h = 0
    
    file.write("pos_charge f fH fxc <H>\n")
    f = np.zeros((pos_charge_grid.size, 7))
    f[:,0] = pos_charge_grid
    
    
    for row, pos_charge in enumerate(pos_charge_grid):
        
        cfg['pos_charge'] = pos_charge
        
        print(f"exact code working on pos_charge = {pos_charge}")
        start = time()
        free_energy = od.freeEnergy(**cfg)
        f[row, 1] = free_energy / pos_charge
        h, mu = od.fctExp(od.energyFct, **cfg)
        f[row, 4] =  h / pos_charge
        f[row, 5] = mu
        f[row, 6] = od.entropy(**cfg)
        t_full += time() - start
        
        print(f"reference code working on pos_charge = {pos_charge}")
        start = time()
        # f[row, 2] -= hr.getHartreeFreeEnergyPerElectron(
        #         cfg['rs'], cfg['theta'], n_spin, cfg['r0'],
        #         pos_charge, cfg['p_max'])
        t_h += time() - start
        
        f[row, 3] = f[row, 1] - f[row, 2]
    
    
        file.write(" ".join("%10.5g" % val for val in f[row]) + "\n")
        
        
    print("full:", t_full)
    print("hartree:", t_h)
    file.close()
    
    
    fig, ax = plt.subplots()
    ax.plot(f[:,0], f[:,1])
    ax.set_xlabel(r"$\cal{N}$")
    # ax.set_ylabel(r"$f_{\mathrm{xc}}$")
    ax.set_ylabel(r"$f$")
    ax.set_title(f"theta = {cfg['theta']}, r0 = {cfg['r0']}, rs = {cfg['rs']}, pmax = {cfg['p_max']}")