#!/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)

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

# 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 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_summary(rms_passed, constraint_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
    
    res_rms = get_status_text(rms_passed)
    res_con = get_status_text(constraint_passed)

    print(f"   [1] Comprehensive RMS check:      {res_rms}")
    print(f"   [2] ADM constraint check:         {res_con}")
    
    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")

    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}")

    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)

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

if __name__ == "__main__":
    main()
