import pandas as pd
import numpy as np
from tqdm import tqdm
from datetime import datetime


def write_cpp_header(df, output_path):
    """Writes the DataFrame to a C++ header as a constexpr array of structs."""
    
    with open(output_path, 'w') as f:
        f.write(f"""// This is an auto generated file, do not edit this file directly
// This file was created using the format.py script in GridFire/utils/WRL
// This file was created on {datetime.now()}
// Rates from this file have been retrived from A. Ravlic, S. Giraud, N. Paar, and R.G.T. Zegers, https://doi.org/10.48550/arXiv.2412.00650
// Rates have been culled to cover T9 < 10 & log electron density < 11 for all species up to and including Z=40 (Zirconium)
// In order to save space in the binary all data has been stored as either single prescision floating point numbers, uint8_t, or uint16_t\n""")
        f.write("#pragma once\n\n")
        f.write("#include <array>\n\n")
        f.write("#include <cstdint>\n\n")
        f.write("namespace gridfire::rates::weak {\n\n")

        f.write("    // Represents a single row from the unified weak rate table\n")
        f.write("    struct RateDataRow {\n")
        f.write("        uint16_t A;\n")
        f.write("        uint8_t Z;\n")
        f.write("        float t9;\n")
        f.write("        float log_rhoye;\n")
        f.write("        float mu_e;\n")
        f.write("        float log_beta_plus;\n")
        f.write("        float log_electron_capture;\n")
        f.write("        float log_neutrino_loss_ec;\n")
        f.write("        float log_beta_minus;\n")
        f.write("        float log_positron_capture;\n")
        f.write("        float log_antineutrino_loss_bd;\n")
        f.write("    };\n\n")

        f.write("    // The complete, pre-processed weak rate table data\n")
        f.write(f"    static constexpr std::array<RateDataRow, {len(df)}> UNIFIED_WEAK_DATA = {{\n")
        
        for row in tqdm(df.itertuples(), total=len(df)):
            f.write("        RateDataRow(")
            f.write(f"{row.A}, {row.Z}, {row.t9}, {row.lrhoYe}, {row.mu_e}, ")
            f.write(f"{row._6}, {row._7}, {row.nue}, {row._9}, {row._10}, {row.nubar}")
            f.write("),\n")
            
        f.write("    };\n\n")
        f.write("} // namespace gridfire::rates::weak\n")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Format unified weak rate library into c++ header")
    parser.add_argument("path", help="path to WRL file", type=str)
    parser.add_argument("-o", "--output", help="path to save c++ header to", default="weak_rate_library.h", type=str)

    args = parser.parse_args()

    print(f"Reading weak rate data from: {args.path}")
    df = pd.read_csv(args.path, sep=r'\s+', comment='#')
    df.columns = [
        "A", "Z", "t9", "lrhoYe", "mu_e", "beta+", "e-",
        "nue", "beta-", "e+", "nubar", "ID"
    ]
    
    df = df[(df.t9 < 10) & (df.lrhoYe < 11) & (df.Z <= 40)]
    print(df)
    print(f"Found {len(df)} total data rows.")
    print(f"Generating C++ header file: {args.output}")
    write_cpp_header(df, args.output)
    print("Done.")