import pandas as pd

# Define fixed-width column specifications based on the format:
# a1 (width 1), i3 (width 3), i5 (width 5), i5 (width 5), i5 (width 5),
# 1x (skip 1), a3 (width 3), a4 (width 4), 1x (skip 1),
# f14.6 (width 14), f12.6 (width 12), f13.5 (width 13),
# 1x (skip 1), f10.5 (width 10), 1x (skip 1),
# a2 (width 2), f13.5 (width 13), f11.5 (width 11),
# 1x (skip 1), i3 (width 3), 1x (skip 1),
# f13.6 (width 13), f12.6 (width 12)
# Compute cumulative positions (0-indexed):
colSpecs = [
    (0, 1),    # control
    (1, 4),    # NZ
    (4, 9),    # N
    (9, 14),   # Z
    (14, 19),  # A
    # skip 1 char at position 19; next field starts at 20
    (20, 23),  # el
    (23, 27),  # o
    # skip 1 char at position 27; next field starts at 28
    (28, 42),  # massExcess (f14.6)
    (42, 54),  # massExcessUnc (f12.6)
    (54, 67),  # bindingEnergy (f13.5)
    # skip 1 char at position 67; next field starts at 68
    (68, 78),  # bindingEnergyUnc (f10.5)
    # skip 1 char at position 78; next field starts at 79
    (79, 81),  # betaCode (a2)
    (81, 94),  # betaDecayEnergy (f13.5)
    (94, 105), # betaDecayEnergyUnc (f11.5)
    # skip 1 char at position 105; next field starts at 106
    (106, 109),# atomicMassInt (i3)
    # skip 1 char at position 109; next field starts at 110
    (110, 123),# atomicMassFrac (f13.6)
    (123, 135) # atomicMassUnc (f12.6)
]

# Define column names (using camelCase for variables)
columnNames = [
    "control",
    "nz",
    "n",
    "z",
    "a",
    "el",
    "o",
    "massExcess",
    "massExcessUnc",
    "bindingEnergy",
    "bindingEnergyUnc",
    "betaCode",
    "betaDecayEnergy",
    "betaDecayEnergyUnc",
    "atomicMassInt",
    "atomicMassFrac",
    "atomicMassUnc"
]

def combine_atomic_mass(row):
    """
    Combine the integer and fractional parts of the atomic mass.
    For example, if atomicMassInt is '1' and atomicMassFrac is '008664.91590',
    this function returns float('1008664.91590').
    """
    intPart = str(row["atomicMassInt"]).strip()
    fracPart = str(row["atomicMassFrac"]).strip()
    try:
        combined = int(intPart) + float(fracPart)/1e6
        return combined
    except ValueError:
        return None

def mkInstanceName(row):
    """
    Make a c++ instance name from the element and atomic number.
    """
    speciesName = f"{row['el'].strip()}-{row['a']}"
    return speciesName.replace("-", "_")

def formatSpecies(row):
    """
    Format c++ instantiation of Species struct from row data.
    """
    name = f"{row['el'].strip()}-{row['a']}"
    instanceName = name.replace("-", "_")
    nz = int(row['nz'])
    n = int(row['n'])
    z = int(row['z'])
    a = int(row['a'])
    bindingEnergy = float(row['bindingEnergy'])
    atomicMass = float(row['atomicMass'])
    atomicMassUnc = float(row['atomicMassUnc'])
    NaN = "std::numeric_limits<double>::quiet_NaN()"
    try:
        betaDecayEnergy = float(row['betaDecayEnergy'].replace("#", "").replace("*", ""))
    except ValueError:
        betaDecayEnergy = NaN
    instantiation = f"static const Species {instanceName}(\"{name}\", \"{row['el']}\", {nz}, {n}, {z}, {a}, {bindingEnergy}, \"{row['betaCode']}\", {betaDecayEnergy}, {atomicMass}, {atomicMassUnc});"
    return instantiation

def formatHeader(dataFrame):
    """
    Format c++ header file from DataFrame.
    """
    header = f"""#ifndef SPECIES_MASS_DATA_H
#define SPECIES_MASS_DATA_H
#include <unordered_map>
#include <string_view>
#include <string>

namespace chemSpecies {{
    struct Species {{
        const std::string_view m_name; //< Name of the species
        const std::string_view m_el; //< Element symbol
        const int m_nz; //< NZ
        const int m_n; //< N
        const int m_z; //< Z
        const int m_a; //< A
        const double m_bindingEnergy; //< Binding energy
        const std::string_view m_betaCode; //< Beta decay code
        const double m_betaDecayEnergy; //< Beta decay energy
        const double m_atomicMass; //< Atomic mass
        const double m_atomicMassUnc; //< Atomic mass uncertainty

        Species(const std::string_view name, const std::string_view el, const int nz, const int n, const int z, const int a, const double bindingEnergy, const std::string_view betaCode, const double betaDecayEnergy, const double atomicMass, const double atomicMassUnc)
            : m_name(name), m_el(el), m_nz(nz), m_n(n), m_z(z), m_a(a), m_bindingEnergy(bindingEnergy), m_betaCode(betaCode), m_betaDecayEnergy(betaDecayEnergy), m_atomicMass(atomicMass), m_atomicMassUnc(atomicMassUnc) {{}};


        double mass() const {{
            return m_atomicMass;
        }}

        double massUnc() const {{
            return m_atomicMassUnc;
        }}

        double bindingEnergy() const {{
            return m_bindingEnergy;
        }}

        double betaDecayEnergy() const {{
            return m_betaDecayEnergy;
        }}

        std::string_view betaCode() const {{
            return m_betaCode;
        }}

        std::string_view name() const {{
            return m_name;
        }}

        std::string_view el() const {{
            return m_el;
        }}

        int nz() const {{
            return m_nz;
        }}

        int n() const {{    
            return m_n;
        }}

        int z() const {{
            return m_z;
        }}

        int a() const {{
            return m_a;
        }}

        friend std::ostream& operator<<(std::ostream& os, const Species& species) {{
            os << static_cast<std::string>(species.m_name) << " (" << species.m_atomicMass << " u)";
            return os;
        }}
    }};
    {'\n    '.join([formatSpecies(row) for index, row in dataFrame.iterrows()])}
    static const std::unordered_map<std::string, Species> species = {{
        {'\n        '.join([f'{{"{row["el"].strip()}-{row["a"]}", {mkInstanceName(row)}}},' for index, row in dataFrame.iterrows()])}
    }};
}}; // namespace chemSpecies
#endif // SPECIES_MASS_DATA_H
"""
    return header

if __name__ == "__main__":
    import argparse 
    import os
    parser = argparse.ArgumentParser(description="Convert mass data to c++ header file.")
    parser.add_argument("input", help="Input file path.")
    parser.add_argument("-o", "--output", help="Output file path.", default="../../assets/static/atomic/include/atomicSpecies.h")
    args = parser.parse_args()


    if not os.path.exists(args.input):
        raise FileNotFoundError(f"File not found: {args.input}")

    # Read the file (adjust the skiprows value if your header differs)
    dataFrame = pd.read_fwf(args.input, colspecs=colSpecs, names=columnNames, skiprows=36)

    # Combine the two atomic mass fields into one float column
    dataFrame["atomicMass"] = dataFrame.apply(combine_atomic_mass, axis=1)
    dataFrame.drop(columns=["atomicMassInt", "atomicMassFrac"], inplace=True)

    # Format the header
    header = formatHeader(dataFrame)
    with open(args.output, "w") as f:
        f.write(header)