import numpy as np
import odeg as od
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.hstack((np.arange(.0625, 12, .0625), np.arange(12, 22, .25)))
    # pos_charge_grid = np.arange(0.25, 8, 0.25)
    pc_left = 1.5
    pc_right = 6.6
    
    pmax_grid = np.arange(1, 6)
    n_spin = 2
    
    cfg = {
        'r0': 0.5,
        'rs': 2,
        'p_max': 5,
        'pos_charge': 10,
        'theta': 5,
        'path': "../results",
        }
    
    fig, axes = plt.subplots(1, 2, figsize = (9, 3))
    
    data = np.zeros((pmax_grid.size, pos_charge_grid.size))
    
    for pmax_ind, pmax in enumerate(pmax_grid):
        red_pc_grid = pos_charge_grid[pos_charge_grid < n_spin * (2*pmax + 1)]
        cfg['p_max'] = pmax
        
        for pc_ind, pos_charge in enumerate(red_pc_grid):
            cfg['pos_charge'] = pos_charge
            
            print(f"exact code working on pmax = {pmax}, pos_charge = {pos_charge}")
            data[pmax_ind, pc_ind] = od.freeEnergy(**cfg) / pos_charge
            
        axes[0].plot(red_pc_grid, data[pmax_ind, :red_pc_grid.size],
                label = r"$p_{\mathrm{max}} = $" + str(pmax))
        
        zoom_ind = (pc_left <= pos_charge_grid) * (pos_charge_grid < min(pc_right, 2*(2*pmax+1)))
        axes[1].plot(pos_charge_grid[zoom_ind], data[pmax_ind, zoom_ind])
        
    # axes[0].set_ylim([-0.63, -0.21])
    axes[0].set_xlabel(r"$\cal{N}$")
    axes[1].set_xlabel(r"$\cal{N}$")
    # ax.set_ylabel(r"$f_{\mathrm{xc}}$")
    axes[0].set_ylabel(r"$F/\cal{N}$" + r"$\,E_{\mathrm{h}}$")
    # ax.set_title(f"theta = {cfg['theta']}, r0 = {cfg['r0']}, rs = {cfg['rs']}")
    
    
    
    fig.legend()
    fig.tight_layout()