import numpy as np
import re
import gzip

EOS2_HEADER = "cccccccccccc"

CE = 1.78266E12
M_U = 931.494

# Index for HShen EOS table, V2
IDX_Log_Baryon_Density_EOS2    = 0
IDX_Baryon_Number_Density_EOS2 = 1
IDX_Proton_Fraction_EOS2       = 2
IDX_Free_Energy_EOS2           = 3
IDX_Internal_Energy_EOS2       = 4
IDX_Entropy_EOS2               = 5
IDX_MU_NP_EOS2                 = 6
IDX_Pressure_EOS2              = 13
IDX_MU_N_EOS2                  = 14
IDX_MU_P_EOS2                  = 15

# EOS2
NB_MESH_EOS2 = np.logspace(5.1, 16.0, num=110, dtype=np.float64)
XP_MESH_EOS2 = np.linspace(0.01, 0.65, num=65, dtype=np.float64)
T_MESH_EOS2  = np.logspace(-1.0, 2.6, num=91, dtype=np.float64)

def skip_line(f):
    f.readline()

def read_numbers(f):
    return [float(x) for x in re.split(r'\s+', f.readline().strip())]

def read_eos(file:str):
    eos_open = gzip.open if file.endswith('.gz') else open
    with eos_open(file) as f:
        for temp in T_MESH_EOS2:
            assert f.readline().strip() == EOS2_HEADER
            skip_line(f)
            (_, temp_r) = read_numbers(f)
            assert np.allclose(temp_r, temp)
            skip_line(f)

            for xp in XP_MESH_EOS2:
                pressures = np.zeros(NB_MESH_EOS2.size)
                for idx, nb in enumerate(NB_MESH_EOS2):
                    pressures[idx] = read_numbers(f)[IDX_Pressure_EOS2]
                    if idx > 0 and nb/CE/M_U >= 0.08 and pressures[idx] < pressures[idx-1]:
                        print(pressures[idx], pressures[idx-1])
                        print("Warning: pressure non-monotonic detected at temp={:.2f}, rho={:.2E} and xp={:.2f}...".format(temp_r, nb/CE/M_U, xp))
                skip_line(f)


if __name__ == "__main__":
    # read_eos("/mnt/c/Users/megrez/Downloads/eos2.tab/eos2.tab")
    read_eos("/path/to/eos2_LU2018_V18.tab")