import os
import sys
import shutil
import subprocess
import numpy as np

def run_small_regression_test():
    workspace_dir = "/home/ignis/workspace/incomp-flame-post"
    code_dir = os.path.join(workspace_dir, "code")
    binary_path = os.path.join(code_dir, "x-edge-cold-bc-uPrime-hybrid")
    
    # 1. Create a temporary testing directory in workspace
    test_run_dir = os.path.join(workspace_dir, "scratch", "small_test_run")
    if os.path.exists(test_run_dir):
        shutil.rmtree(test_run_dir)
    os.makedirs(test_run_dir, exist_ok=True)
    
    # 2. Copy config files from historical backup 'ic1'
    backup_ic1_dir = "/tank/ingest/from_sata_sdk_4tb/161/ignis/dns/incomp-flame-post/ic1"
    
    shutil.copy(os.path.join(backup_ic1_dir, "post-edge-cold-bc-hybrid-intro"), test_run_dir)
    shutil.copy(os.path.join(backup_ic1_dir, "otape"), test_run_dir)
    shutil.copy(binary_path, test_run_dir)
    
    # Update config to run only 6 files (1000 to 1005)
    intro_path = os.path.join(test_run_dir, "post-edge-cold-bc-hybrid-intro")
    with open(intro_path, "r") as f:
        lines = f.readlines()
    with open(intro_path, "w") as f:
        for line in lines:
            if line.strip().startswith("endnum"):
                f.write("    endnum 1005        ! only 6 files for fast test\n")
            else:
                f.write(line)
                
    # 3. Create symlinks for fort.1000 to fort.1005
    print("Linking subset of fort.* data files...")
    linked_count = 0
    for i in range(1000, 1006):
        src_file = os.path.join(backup_ic1_dir, f"fort.{i}")
        dst_link = os.path.join(test_run_dir, f"fort.{i}")
        os.symlink(src_file, dst_link)
        linked_count += 1
            
    print(f"Linked {linked_count} data files to temporary run directory.")
    
    # 4. Run the compiled binary using mpirun
    print("Running compiled MPI binary inside test run directory...")
    # Using 4 ranks for the small test run (under 5 seconds)
    run_cmd = ["mpirun", "--oversubscribe", "-np", "4", "./x-edge-cold-bc-uPrime-hybrid"]
    try:
        result = subprocess.run(
            run_cmd,
            cwd=test_run_dir,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=True
        )
        print("MPI Binary run completed successfully!")
    except subprocess.CalledProcessError as e:
        print("MPI Binary execution FAILED!")
        sys.exit(1)
        
    # 5. Load and compare qEdge_X.dat
    new_result_path = os.path.join(test_run_dir, "qEdge_X.dat")
    ref_result_path = os.path.join(code_dir, "code_gen", "qEdge_X_small_ref.dat")
    
    if not os.path.exists(new_result_path):
        print(f"Error: Generated qEdge_X.dat not found in {test_run_dir}!")
        sys.exit(1)
        
    print("Comparing new qEdge_X.dat with reference qEdge_X_small_ref.dat...")
    
    # Parse headers
    with open(new_result_path, 'r') as f:
        new_header = f.readline().strip().split()
    with open(ref_result_path, 'r') as f:
        ref_header = f.readline().strip().split()
        
    # Read numeric data
    new_data = np.loadtxt(new_result_path, skiprows=1)
    ref_data = np.loadtxt(ref_result_path, skiprows=1)
    
    if new_data.shape != ref_data.shape:
        print(f"Shape Mismatch: New data shape {new_data.shape} != Reference data shape {ref_data.shape}")
        sys.exit(1)
        
    header_mapping = {}
    for i, col_name in enumerate(new_header):
        if col_name in ref_header:
            header_mapping[i] = ref_header.index(col_name)
            
    tolerance = 2e-10
    max_diffs = []
    failed = False
    
    print("\nColumn-by-column numeric verification:")
    print(f"{'Variable Name':<20} | {'Max Abs Error':<13} | {'Max Rel Error':<13} | {'Status':<10}")
    print("-" * 68)
    
    for new_idx, ref_idx in header_mapping.items():
        var_name = new_header[new_idx]
        new_col = new_data[:, new_idx]
        ref_col = ref_data[:, ref_idx]
        
        nan_mismatch = not np.array_equal(np.isnan(new_col), np.isnan(ref_col))
        
        abs_diff = np.abs(new_col - ref_col)
        if np.all(np.isnan(abs_diff)):
            max_diff = 0.0
        else:
            max_diff = np.nanmax(abs_diff)
        
        max_diffs.append(max_diff)
        
        with np.errstate(divide='ignore', invalid='ignore'):
            rel_diff = abs_diff / np.abs(ref_col)
            rel_diff[np.isnan(rel_diff) | np.isinf(rel_diff) | (np.abs(ref_col) < 1e-12)] = 0.0
            if np.all(rel_diff == 0.0):
                max_rel = 0.0
            else:
                max_rel = np.nanmax(rel_diff)
        
        status = "PASSED"
        if nan_mismatch:
            status = "NAN_MISMATCH"
            failed = True
        else:
            is_passed = (max_diff <= tolerance) or (max_rel <= 1e-4)
            if not is_passed:
                status = "FAILED"
                failed = True
            
        print(f"{var_name:<20} | {max_diff:<13.4e} | {max_rel:<13.4e} | {status:<10}")
        
    if failed:
        print("\nSmall Regression Verification FAILED!")
        sys.exit(1)
    else:
        print(f"\nSmall Regression Verification SUCCESSFUL! All columns matched within tolerance.")
        shutil.rmtree(test_run_dir) # cleanup temp test run directory upon success
        sys.exit(0)

if __name__ == '__main__':
    run_small_regression_test()