#!/usr/bin/env python3
"""
AMSS-NCKU GW150914 Simulation Regression Test Script (Comprehensive Version)

Verification Requirements:
1. RMS errors < 1% for:
   - 3D Vector Total RMS
   - X Component RMS
   - Y Component RMS
   - Z Component RMS
2. ADM constraint violation < 2 (Grid Level 0)
3. The following figure PDFs must match GW150914-origin exactly after rasterization:
   - ADM_Constraint_Grid_Level_0.pdf
   - BH_Trajectory_21_XY.pdf
   - BH_Trajectory_XY.pdf
   The script also reports the percentage of differing pixels for each figure.

RMS Calculation Method:
- Computes trajectory deviation on the XY plane independently for BH1 and BH2
- For each black hole: RMS = sqrt((1/M) * sum((Δr_i / r_i^max)^2)) × 100%
- Final RMS = max(RMS_BH1, RMS_BH2)

Usage: python3 AMSS_NCKU_Verify_ASC26.py [output_dir]
Default: output_dir = GW150914/AMSS_NCKU_output
Reference: GW150914-origin (baseline simulation)
"""

import numpy as np
import sys
import os
import shutil
import subprocess
import tempfile
from PIL import Image

# ANSI Color Codes
class Color:
    GREEN = '\033[92m'
    RED = '\033[91m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    BOLD = '\033[1m'
    RESET = '\033[0m'

def get_status_text(passed):
    if passed:
        return f"{Color.GREEN}{Color.BOLD}PASS{Color.RESET}"
    else:
        return f"{Color.RED}{Color.BOLD}FAIL{Color.RESET}"

def load_bh_trajectory(filepath):
    """Load black hole trajectory data"""
    data = np.loadtxt(filepath)
    return {
        'time': data[:, 0],
        'x1': data[:, 1], 'y1': data[:, 2], 'z1': data[:, 3],
        'x2': data[:, 4], 'y2': data[:, 5], 'z2': data[:, 6]
    }


def load_constraint_data(filepath):
    """Load constraint violation data"""
    data = []
    with open(filepath, 'r') as f:
        for line in f:
            if line.startswith('#'):
                continue
            parts = line.split()
            if len(parts) >= 8:
                data.append([float(x) for x in parts[:8]])
    return np.array(data)


def resolve_figure_dir(path):
    """Resolve the sibling figure directory from an output or figure path."""
    normalized = os.path.normpath(path)
    if os.path.basename(normalized) == "figure":
        return normalized
    return os.path.join(os.path.dirname(normalized), "figure")


def render_pdf_to_images(pdf_path, dpi=150):
    """Render a PDF to RGB images using Ghostscript."""
    gs_path = shutil.which("gs")
    if gs_path is None:
        raise RuntimeError("Ghostscript executable 'gs' was not found in PATH")

    with tempfile.TemporaryDirectory(prefix="amss_verify_pdf_") as temp_dir:
        output_pattern = os.path.join(temp_dir, "page-%03d.ppm")
        cmd = [
            gs_path,
            "-q",
            "-dSAFER",
            "-dBATCH",
            "-dNOPAUSE",
            "-sDEVICE=ppmraw",
            f"-r{dpi}",
            f"-o{output_pattern}",
            pdf_path
        ]

        try:
            subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
        except subprocess.CalledProcessError as exc:
            message = exc.stderr.strip() or str(exc)
            raise RuntimeError(f"Failed to render PDF '{pdf_path}': {message}") from exc

        ppm_files = sorted(
            os.path.join(temp_dir, filename)
            for filename in os.listdir(temp_dir)
            if filename.endswith(".ppm")
        )

        if not ppm_files:
            raise RuntimeError(f"No rendered pages were produced for '{pdf_path}'")

        images = []
        for ppm_file in ppm_files:
            with Image.open(ppm_file) as img:
                images.append(np.array(img.convert("RGB"), dtype=np.uint8))

        return images


def compare_rendered_pages(ref_img, target_img):
    """Return (different_pixels, total_pixels) for two rendered RGB pages."""
    ref_h, ref_w = ref_img.shape[:2]
    tgt_h, tgt_w = target_img.shape[:2]
    total_pixels = max(ref_h, tgt_h) * max(ref_w, tgt_w)

    if ref_h == tgt_h and ref_w == tgt_w:
        different_pixels = int(np.count_nonzero(np.any(ref_img != target_img, axis=2)))
        return different_pixels, total_pixels

    diff_mask = np.ones((max(ref_h, tgt_h), max(ref_w, tgt_w)), dtype=bool)
    overlap_h = min(ref_h, tgt_h)
    overlap_w = min(ref_w, tgt_w)
    overlap_diff = np.any(ref_img[:overlap_h, :overlap_w] != target_img[:overlap_h, :overlap_w], axis=2)
    diff_mask[:overlap_h, :overlap_w] = overlap_diff
    different_pixels = int(np.count_nonzero(diff_mask))
    return different_pixels, total_pixels


def compare_pdf_images(ref_pdf, target_pdf, dpi=150, threshold_percent=0.001):
    """Compare two PDFs by rasterizing them and counting differing pixels."""
    ref_pages = render_pdf_to_images(ref_pdf, dpi=dpi)
    target_pages = render_pdf_to_images(target_pdf, dpi=dpi)

    total_pixels = 0
    different_pixels = 0
    max_pages = max(len(ref_pages), len(target_pages))

    for page_idx in range(max_pages):
        if page_idx < len(ref_pages) and page_idx < len(target_pages):
            page_diff, page_total = compare_rendered_pages(ref_pages[page_idx], target_pages[page_idx])
        else:
            existing_page = ref_pages[page_idx] if page_idx < len(ref_pages) else target_pages[page_idx]
            page_total = existing_page.shape[0] * existing_page.shape[1]
            page_diff = page_total

        total_pixels += page_total
        different_pixels += page_diff

    diff_percent = (different_pixels / total_pixels * 100.0) if total_pixels else 0.0
    return {
        "different_pixels": different_pixels,
        "total_pixels": total_pixels,
        "diff_percent": diff_percent,
        "pages_ref": len(ref_pages),
        "pages_target": len(target_pages),
        "passed": diff_percent < threshold_percent
    }


def compare_required_figures(reference_figure_dir, target_figure_dir):
    """Compare the required GW150914 figure PDFs."""
    figure_names = [
        "ADM_Constraint_Grid_Level_0.pdf",
        "BH_Trajectory_21_XY.pdf",
        "BH_Trajectory_XY.pdf"
    ]

    results = []
    for figure_name in figure_names:
        ref_pdf = os.path.join(reference_figure_dir, figure_name)
        target_pdf = os.path.join(target_figure_dir, figure_name)

        if not os.path.exists(ref_pdf):
            raise FileNotFoundError(f"Reference figure not found: {ref_pdf}")
        if not os.path.exists(target_pdf):
            raise FileNotFoundError(f"Target figure not found: {target_pdf}")

        comparison = compare_pdf_images(ref_pdf, target_pdf)
        comparison["name"] = figure_name
        results.append(comparison)

    return results

def calculate_all_rms_errors(bh_data_ref, bh_data_target):
    """
    Calculate 3D Vector RMS and component-wise RMS (X, Y, Z) independently.
    Uses r = sqrt(x^2 + y^2) as the denominator for all error normalizations.
    Returns the maximum error between BH1 and BH2 for each category.
    """
    M = min(len(bh_data_ref['time']), len(bh_data_target['time']))

    if M < 10:
        return None, "Insufficient data points for comparison"

    results = {}

    for bh in ['1', '2']:
        x_r, y_r, z_r = bh_data_ref[f'x{bh}'][:M], bh_data_ref[f'y{bh}'][:M], bh_data_ref[f'z{bh}'][:M]
        x_n, y_n, z_n = bh_data_target[f'x{bh}'][:M], bh_data_target[f'y{bh}'][:M], bh_data_target[f'z{bh}'][:M]

        # 核心修改：根据组委会的邮件指示，分母统一使用 r = sqrt(x^2 + y^2)
        r_ref = np.sqrt(x_r**2 + y_r**2)
        r_new = np.sqrt(x_n**2 + y_n**2)
        denom_max = np.maximum(r_ref, r_new)

        valid = denom_max > 1e-15
        if np.sum(valid) < 10:
            results[f'BH{bh}'] = { '3D_Vector': 0.0, 'X_Component': 0.0, 'Y_Component': 0.0, 'Z_Component': 0.0 }
            continue

        def calc_rms(delta):
            # 将对应分量的偏差除以统一的轨道半径分母 denom_max
            return np.sqrt(np.mean((delta[valid] / denom_max[valid])**2)) * 100

        # 1. Total 3D Vector RMS
        delta_vec = np.sqrt((x_r - x_n)**2 + (y_r - y_n)**2 + (z_r - z_n)**2)
        rms_3d = calc_rms(delta_vec)

        # 2. Component-wise RMS (分离计算各轴，但共用半径分母)
        rms_x = calc_rms(np.abs(x_r - x_n))
        rms_y = calc_rms(np.abs(y_r - y_n))
        rms_z = calc_rms(np.abs(z_r - z_n))

        results[f'BH{bh}'] = {
            '3D_Vector': rms_3d,
            'X_Component': rms_x,
            'Y_Component': rms_y,
            'Z_Component': rms_z
        }

    # 获取 BH1 和 BH2 中的最大误差
    max_rms = {
        '3D_Vector': max(results['BH1']['3D_Vector'], results['BH2']['3D_Vector']),
        'X_Component': max(results['BH1']['X_Component'], results['BH2']['X_Component']),
        'Y_Component': max(results['BH1']['Y_Component'], results['BH2']['Y_Component']),
        'Z_Component': max(results['BH1']['Z_Component'], results['BH2']['Z_Component'])
    }

    return max_rms, None

def analyze_constraint_violation(constraint_data, n_levels=9):
    """
    Analyze ADM constraint violation
    Return maximum constraint violation for Grid Level 0
    """
    # Extract Grid Level 0 data (first entry for each time step)
    level0_data = constraint_data[::n_levels]

    # Calculate maximum absolute value for each constraint
    results = {
        'Ham': np.max(np.abs(level0_data[:, 1])),
        'Px': np.max(np.abs(level0_data[:, 2])),
        'Py': np.max(np.abs(level0_data[:, 3])),
        'Pz': np.max(np.abs(level0_data[:, 4])),
        'Gx': np.max(np.abs(level0_data[:, 5])),
        'Gy': np.max(np.abs(level0_data[:, 6])),
        'Gz': np.max(np.abs(level0_data[:, 7]))
    }

    results['max_violation'] = max(results.values())
    return results


def print_header():
    print("\n" + Color.BLUE + Color.BOLD + "=" * 65 + Color.RESET)
    print(Color.BOLD + "   AMSS-NCKU GW150914 Comprehensive Regression Test" + Color.RESET)
    print(Color.BLUE + Color.BOLD + "=" * 65 + Color.RESET)

def print_rms_results(rms_dict, error, threshold=1.0):
    print(f"\n{Color.BOLD}1. RMS Error Analysis (Maximums of BH1 & BH2){Color.RESET}")
    print("-" * 65)

    if error:
        print(f"   {Color.RED}Error: {error}{Color.RESET}")
        return False

    all_passed = True
    print(f"   Requirement: < {threshold}%\n")

    for key, val in rms_dict.items():
        passed = val < threshold
        all_passed = all_passed and passed
        status = get_status_text(passed)
        print(f"   {key:15}: {val:8.4f}%   |   Status: {status}")

    return all_passed

def print_constraint_results(results, threshold=2.0):
    print(f"\n{Color.BOLD}2. ADM Constraint Violation Analysis (Grid Level 0){Color.RESET}")
    print("-" * 65)

    names = ['Ham', 'Px', 'Py', 'Pz', 'Gx', 'Gy', 'Gz']
    for i, name in enumerate(names):
        print(f"   Max |{name:3}|: {results[name]:.6f}", end="   ")
        if (i + 1) % 2 == 0: print()
    if len(names) % 2 != 0: print()

    passed = results['max_violation'] < threshold
    
    print(f"\n   Maximum violation:  {results['max_violation']:.6f}")
    print(f"   Requirement:        < {threshold}")
    print(f"   Status:             {get_status_text(passed)}")

    return passed


def print_figure_results(results, threshold_percent=0.001):
    print(f"\n{Color.BOLD}3. Figure Pixel Comparison (PDF Rasterization){Color.RESET}")
    print("-" * 65)
    print(f"   Requirement: < {threshold_percent:.3f}% differing pixels\n")

    all_passed = True
    for result in results:
        passed = result["passed"]
        all_passed = all_passed and passed
        status = get_status_text(passed)
        print(f"   {result['name']:32}: {result['diff_percent']:10.6f}%   |   Status: {status}")

        if result["pages_ref"] != result["pages_target"]:
            print(f"   {'':32}  pages(ref/target): {result['pages_ref']}/{result['pages_target']}")

    return all_passed


def print_figure_error(error_message):
    print(f"\n{Color.BOLD}3. Figure Pixel Comparison (PDF Rasterization){Color.RESET}")
    print("-" * 65)
    print(f"   {Color.RED}Error: {error_message}{Color.RESET}")
    return False


def print_summary(rms_passed, constraint_passed, figure_passed):
    print("\n" + Color.BLUE + Color.BOLD + "=" * 65 + Color.RESET)
    print(Color.BOLD + "Verification Summary" + Color.RESET)
    print(Color.BLUE + Color.BOLD + "=" * 65 + Color.RESET)

    all_passed = rms_passed and constraint_passed and figure_passed
    
    res_rms = get_status_text(rms_passed)
    res_con = get_status_text(constraint_passed)
    res_fig = get_status_text(figure_passed)

    print(f"   [1] Comprehensive RMS check:      {res_rms}")
    print(f"   [2] ADM constraint check:         {res_con}")
    print(f"   [3] Figure pixel comparison:      {res_fig}")
    
    final_status = f"{Color.GREEN}{Color.BOLD}ALL CHECKS PASSED{Color.RESET}" if all_passed else f"{Color.RED}{Color.BOLD}SOME CHECKS FAILED{Color.RESET}"
    print(f"\n   Overall result: {final_status}")
    print(Color.BLUE + Color.BOLD + "=" * 65 + Color.RESET + "\n")

    return all_passed

def main():
    if len(sys.argv) > 1:
        target_dir = sys.argv[1]
    else:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        target_dir = os.path.join(script_dir, "GW150914/AMSS_NCKU_output")

    script_dir = os.path.dirname(os.path.abspath(__file__))
    reference_dir = os.path.join(script_dir, "GW150914-origin/AMSS_NCKU_output")
    target_figure_dir = resolve_figure_dir(target_dir)
    reference_figure_dir = os.path.join(script_dir, "GW150914-origin/figure")

    bh_file_ref = os.path.join(reference_dir, "bssn_BH.dat")
    bh_file_target = os.path.join(target_dir, "bssn_BH.dat")
    constraint_file = os.path.join(target_dir, "bssn_constraint.dat")

    if not os.path.exists(bh_file_ref):
        print(f"{Color.RED}{Color.BOLD}Error:{Color.RESET} Baseline trajectory file not found: {bh_file_ref}")
        sys.exit(1)
    if not os.path.exists(bh_file_target):
        print(f"{Color.RED}{Color.BOLD}Error:{Color.RESET} Target trajectory file not found: {bh_file_target}")
        sys.exit(1)
    if not os.path.exists(constraint_file):
        print(f"{Color.RED}{Color.BOLD}Error:{Color.RESET} Constraint data file not found: {constraint_file}")
        sys.exit(1)

    print_header()
    print(f"\n{Color.BOLD}Reference (Baseline):{Color.RESET} {Color.BLUE}{reference_dir}{Color.RESET}")
    print(f"{Color.BOLD}Target (Optimized):  {Color.RESET} {Color.BLUE}{target_dir}{Color.RESET}")
    print(f"{Color.BOLD}Reference Figures:   {Color.RESET} {Color.BLUE}{reference_figure_dir}{Color.RESET}")
    print(f"{Color.BOLD}Target Figures:      {Color.RESET} {Color.BLUE}{target_figure_dir}{Color.RESET}")

    bh_data_ref = load_bh_trajectory(bh_file_ref)
    bh_data_target = load_bh_trajectory(bh_file_target)
    constraint_data = load_constraint_data(constraint_file)

    # Output modified RMS results
    rms_dict, error = calculate_all_rms_errors(bh_data_ref, bh_data_target)
    rms_passed = print_rms_results(rms_dict, error)

    # Output constraint results
    constraint_results = analyze_constraint_violation(constraint_data)
    constraint_passed = print_constraint_results(constraint_results)

    try:
        figure_results = compare_required_figures(reference_figure_dir, target_figure_dir)
        figure_passed = print_figure_results(figure_results)
    except (FileNotFoundError, RuntimeError) as exc:
        figure_passed = print_figure_error(str(exc))

    all_passed = print_summary(rms_passed, constraint_passed, figure_passed)
    sys.exit(0 if all_passed else 1)

if __name__ == "__main__":
    main()
