import re
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import collections

def main():
    data_file = 'rawdata-kcurve2_3_108-112-24-56.txt'
    output_image = 'gr_mass_ratio_analysis.png'

    print(f"Parsing data from {data_file}...")
    try:
        with open(data_file, 'r') as f:
            text = f.read()
    except FileNotFoundError:
        print(f"Error: {data_file} not found.")
        return

    datasets = text.split('Initial Parameters:')
    parsed_data = collections.defaultdict(list)

    for ds in datasets[1:]:
        m_kcurve = re.search(r'kcurve:\s*([\d\.]+)', ds)
        if not m_kcurve: continue
        kcurve = m_kcurve.group(1)
        
        m_points = re.search(r'points:\s*(\d+)', ds)
        if not m_points: continue
        M = int(m_points.group(1)) - 1
        
        # Priority: Velocity comparison, fallback to generic 1 comparison
        m_gr = re.search(r'GR comparison velocity \.\.\.\s*([\d\.]+)', ds)
        if not m_gr:
            m_gr = re.search(r'GR comparison 1 \.\.\.\s*([\d\.]+)', ds)
            
        if not m_gr: 
            print(f"Skipping dataset kcurve={kcurve}, M={M} (No GR comparison found)")
            continue
            
        gr_ratio = float(m_gr.group(1))
        parsed_data[kcurve].append((M, gr_ratio))
        print(f"Parsed: kcurve={kcurve:<5} M={M:<5} GR_Ratio={gr_ratio:.4f}")

    # Plot Setup
    plt.style.use('dark_background')
    fig, ax = plt.subplots(figsize=(12, 8), facecolor='#0b0b0b')
    ax.set_facecolor('#0b0b0b')

    colors = {'2': '#ff0055', '3': '#ff9900', '108': '#00d4ff', '112': '#00ff88'}
    labels = {'2': 'High Eccentricity (e=0.93, kcurve=2)', 
              '3': 'Med Eccentricity (e=0.85, kcurve=3)', 
              '108': 'Low Eccentricity (e=0.20, kcurve=108)', 
              '112': 'Low Eccentricity (e=0.20, kcurve=112)'}

    for kcurve, points in parsed_data.items():
        points.sort(key=lambda x: x[0])
        M_vals = [p[0] for p in points]
        gr_vals = [p[1] for p in points]
        
        ax.plot(M_vals, gr_vals, marker='o', linewidth=3, markersize=8, 
                color=colors.get(kcurve, 'white'), label=labels.get(kcurve, f'kcurve={kcurve}'))

    ax.set_title("General Relativity Precession Amplification vs. Mass Ratio", fontsize=18, color='white', pad=20)
    ax.set_xlabel("Central Mass M (Planck Masses)", fontsize=14, color='#aaaaaa')
    ax.set_ylabel("GR Amplification Ratio (Simulation / Standard GR Prediction)", fontsize=14, color='#aaaaaa')
    
    # Linear scale requested
    ax.set_yscale('linear')
    ax.grid(True, color='#333333', linestyle='--', alpha=0.5)

    # Reference Line
    ax.axhline(y=1.0, color='#888888', linestyle=':', linewidth=2, label="Standard Macroscopic GR Limit (Ratio = 1.0)")

    ax.legend(fontsize=12, frameon=True, facecolor='#111111', loc='upper right')

    plt.tight_layout()
    plt.savefig(output_image, dpi=150, facecolor='#0b0b0b')
    print(f"\nPlot successfully generated: {output_image}")

if __name__ == "__main__":
    main()