import numpy as np
from ezfio import ezfio
from datetime import datetime
import time
from matplotlib import pyplot as plt
from scipy.sparse import coo_matrix



def plot_coo_matrix(m):
    if not isinstance(m, coo_matrix):
        m = coo_matrix(m)
    fig = plt.figure()
    ax = fig.add_subplot(111, facecolor='black')
    ax.plot(m.col, m.row, 's', color='white', ms=1)
    ax.set_xlim(0, m.shape[1])
    ax.set_ylim(0, m.shape[0])
    ax.set_aspect('equal')
    for spine in ax.spines.values():
        spine.set_visible(False)
    ax.invert_yaxis()
    ax.set_aspect('equal')
    ax.set_xticks([])
    ax.set_yticks([])
    return ax


if __name__ == "__main__":

    t0 = time.time()

    EZFIO_file = "/home/aammar/qp2/src/svdwf/h2o_work/FN_test/cc_pCVDZ/h2o_dz"
    ezfio.set_file(EZFIO_file)
    print(" EZFIO = {}\n".format(EZFIO_file))

    n_alpha = ezfio.get_spindeterminants_n_det_alpha()
    n_beta  = ezfio.get_spindeterminants_n_det_beta()
    print(' matrix: {} x {}\n'.format(n_alpha,n_beta))

    n_selected = 15 * 15

    # selected with v2
    f     = open("fort.222", "r")
    lines = f.readlines()
    f.close()
    E_v2      = float(lines[-1].split()[4])
    deltaE_v2 = np.zeros((n_alpha,n_beta))
    rows, cols, vals = [], [], []
    print("E_v2 = {}\n".format(E_v2))
    for i in range(n_selected):
        line   = lines[i].split()
        ia, ib = int(line[1]), int(line[2])
        deltaE_v2[ia,ib] = float(line[3])/E_v2
        rows.append(ia)
        cols.append(ib)
        vals.append(deltaE_v2[ia,ib])

    m   = coo_matrix((vals, (rows, cols)), shape=(n_alpha,n_beta))
    ax  = plot_coo_matrix(m)
    ax.figure.savefig("selected_SVD.pdf")