#!/usr/bin/env python3
import os
import sys
import argparse
import struct
import subprocess
import time
import re

def parse_args():
    parser = argparse.ArgumentParser(description="Auto-tune transpose block size for m_calculate.f90")
    parser.add_argument("--nxp", type=int, help="Override grid size in X direction")
    parser.add_argument("--nyp", type=int, help="Override grid size in Y direction")
    parser.add_argument("--nzp", type=int, help="Override grid size in Z direction")
    parser.add_argument("--iters", type=int, default=5, help="Number of iterations for benchmark run")
    parser.add_argument("--apply", action="store_true", help="Apply the optimal block size to the makefile")
    return parser.parse_args()

def read_fortran_grid_sizes(filepath):
    try:
        with open(filepath, 'rb') as f:
            # First 4 bytes is the record length
            len_bytes = f.read(4)
            if len(len_bytes) < 4:
                return None
            rec_len = struct.unpack('i', len_bytes)[0]
            
            # Read the record data
            rec_data = f.read(rec_len)
            if len(rec_data) < rec_len:
                return None
            
            # Record structure: tnow (double), nx (int64), ny (int64), nz (int64)
            # Since compiler has -fdefault-integer-8, these are 64-bit integers (8 bytes)
            # Format: 'dqqq' -> double, int64, int64, int64
            tnow, nx, ny, nz = struct.unpack('dqqq', rec_data[:32])
            return nx, ny, nz
    except Exception as e:
        print(f"Error reading grid size from {filepath}: {e}")
        return None

def detect_grid_sizes():
    # Try current directory first, then parent directory
    search_dirs = [".", "..", "../scratch/ic1_regression_test", "scratch/ic1_regression_test"]
    
    intro_file = None
    intro_dir = None
    for d in search_dirs:
        p = os.path.join(d, "post-edge-cold-bc-hybrid-intro")
        if os.path.exists(p):
            intro_file = p
            intro_dir = d
            break
            
    if intro_file:
        print(f"Found intro config file at {intro_file}")
        startnum = None
        try:
            with open(intro_file, 'r') as f:
                for line in f:
                    if 'startnum' in line:
                        # Extract the first integer in the line
                        match = re.search(r'\d+', line)
                        if match:
                            startnum = int(match.group())
                            break
        except Exception as e:
            print(f"Error reading startnum from {intro_file}: {e}")
            
        if startnum is not None:
            # Search for fort.<startnum> in search_dirs
            for d in search_dirs:
                fort_path = os.path.join(d, f"fort.{startnum}")
                if os.path.exists(fort_path):
                    sizes = read_fortran_grid_sizes(fort_path)
                    if sizes:
                        print(f"Successfully auto-detected grid size from {fort_path}: NX={sizes[0]}, NY={sizes[1]}, NZ={sizes[2]}")
                        return sizes
    
    # Fallback to default
    print("Could not auto-detect grid sizes. Using defaults: NX=512, NY=256, NZ=256")
    return 512, 256, 256

def find_divisors(n):
    divisors = []
    for i in range(1, int(n**0.5) + 1):
        if n % i == 0:
            divisors.append(i)
            if i*i != n:
                divisors.append(n // i)
    return sorted(divisors)

def run_command(args, cwd=None):
    res = subprocess.run(args, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    return res.returncode == 0, res.stdout, res.stderr

def update_makefile(best_bs):
    makefile_path = "makefile"
    if not os.path.exists(makefile_path):
        makefile_path = "code/makefile"
    if not os.path.exists(makefile_path):
        makefile_path = "../code/makefile"
    if not os.path.exists(makefile_path):
        print("Error: makefile not found. Cannot apply changes.")
        return False
        
    try:
        with open(makefile_path, 'r') as f:
            content = f.read()
            
        new_content = re.sub(r'(BLOCKSIZE\s*\?=\s*)\d+', f'\\g<1>{best_bs}', content)
        
        with open(makefile_path, 'w') as f:
            f.write(new_content)
            
        print(f"Successfully updated {makefile_path} with BLOCKSIZE={best_bs}")
        return True
    except Exception as e:
        print(f"Error updating makefile: {e}")
        return False

def main():
    args = parse_args()
    
    # Grid sizes detection
    detected_nx, detected_ny, detected_nz = detect_grid_sizes()
    nx = args.nxp if args.nxp else detected_nx
    ny = args.nyp if args.nyp else detected_ny
    nz = args.nzp if args.nzp else detected_nz
    
    print(f"Target grid size for tuning: NX={nx}, NY={ny}, NZ={nz}")
    
    # Calculate candidate block sizes (must be divisors of NY)
    all_divisors = find_divisors(ny)
    # Filter to block sizes >= 4 to avoid excessive overhead
    candidates = [d for d in all_divisors if d >= 4]
    
    if not candidates:
        print(f"No suitable divisors found for NY={ny}. Testing all divisors: {all_divisors}")
        candidates = all_divisors
        
    print(f"Candidate block sizes to test: {candidates}")
    
    code_dir = "."
    if not os.path.exists("makefile") and os.path.exists("code/makefile"):
        code_dir = "code"
        
    print(f"Using build directory: {os.path.abspath(code_dir)}")
    print(f"Running each benchmark with {args.iters} iterations...")
    print("\n" + "="*80)
    print(f"{'Block Size':^12} | {'Compile Time (s)':^18} | {'Avg Run Time (s)':^20} | {'Status':^12}")
    print("-"*80)
    
    results = []
    
    for bs in candidates:
        # 1. Compile
        t_start = time.time()
        ok_clean, _, _ = run_command(["make", "clean"], cwd=code_dir)
        if not ok_clean:
            print(f"{bs:^12} | {'-':^18} | {'-':^20} | {'CLEAN FAILED':^12}")
            continue
            
        compile_cmd = ["make", "benchmark_transpose", f"BLOCKSIZE={bs}"]
        ok_compile, stdout_c, stderr_c = run_command(compile_cmd, cwd=code_dir)
        t_compile = time.time() - t_start
        
        if not ok_compile:
            print(f"{bs:^12} | {t_compile:^18.2f} | {'-':^20} | {'COMP. FAILED':^12}")
            # print(stderr_c)
            continue
            
        # 2. Run benchmark
        # benchmark_transpose binary is in code_dir
        binary_path = "./benchmark_transpose"
        run_args = [binary_path, str(nx), str(ny), str(nz), str(args.iters)]
        
        ok_run, stdout_r, stderr_r = run_command(run_args, cwd=code_dir)
        if not ok_run:
            print(f"{bs:^12} | {t_compile:^18.2f} | {'-':^20} | {'RUN FAILED':^12}")
            # print(stderr_r)
            continue
            
        try:
            # Parse the time output from the program
            avg_time = float(stdout_r.strip())
            print(f"{bs:^12d} | {t_compile:^18.2f} | {avg_time:^20.6f} | {'SUCCESS':^12}")
            results.append((bs, avg_time))
        except ValueError:
            print(f"{bs:^12} | {t_compile:^18.2f} | {'-':^20} | {'PARSE ERROR':^12}")
            # print(f"Output was: {stdout_r}")
            
    print("="*80)
    
    if not results:
        print("Tuning failed: no successful candidate runs.")
        sys.exit(1)
        
    # Find the best block size
    best_bs, best_time = min(results, key=lambda x: x[1])
    print(f"\nOptimal BLOCKSIZE found: {best_bs} with average time of {best_time:.6f} seconds.")
    
    # Recommendations
    print("\nTested results:")
    for bs, t in results:
        rel = (t - best_time) / best_time * 100.0
        marker = " (Best)" if bs == best_bs else f" (+{rel:.1f}%)"
        print(f"  Block Size {bs:3d}: {t:.6f} s{marker}")
        
    if args.apply:
        update_makefile(best_bs)
    else:
        print(f"\nTo apply this block size, run with --apply or manually update 'BLOCKSIZE?={best_bs}' in the makefile.")

if __name__ == '__main__':
    main()