import pyvista as pv
import numpy as np
from scipy.spatial.distance import jensenshannon
import os


dataset = "plane_transonic"
train_folder = f'/raid/ansysai/pkakka/6-Transformers/comparePhysicsLM/Data/{dataset}/'
train_save_dir = train_folder + f"../../metrics/{dataset}/"
test_dir = train_save_dir + f"Transformer_baseline_ep1000/best_case/vtk_files/"

# Function to extract pressure values from VTP file
# For training data: uses 'pressure' field
# For test data: can extract both 'gt_pressure' and 'pred_pressure'
def get_pressures(vtp_file, pressure_field='pressure'):
    mesh = pv.read(vtp_file)
    if pressure_field not in mesh.point_data:
        raise ValueError(f"No '{pressure_field}' data in {vtp_file}")
    return mesh.point_data[pressure_field]

# Function to get both GT and predicted pressures from test VTK file
def get_both_pressures(vtk_file):
    mesh = pv.read(vtk_file)
    
    gt_pressure = None
    pred_pressure = None
    
    if 'gt_pressure' in mesh.point_data:
        gt_pressure = mesh.point_data['gt_pressure']
    elif 'pressure' in mesh.point_data:
        gt_pressure = mesh.point_data['pressure']
    
    if 'pred_pressure' in mesh.point_data:
        pred_pressure = mesh.point_data['pred_pressure']
    
    return gt_pressure, pred_pressure

# Part 1: Compute and store training PDF (histogram)
def compute_training_pdf(train_folder, output_file='train_dist.npz', num_bins=50):
    # List all VTP files in training folder
    train_files = [os.path.join(train_folder, f) for f in os.listdir(train_folder) if f.endswith('.vtp')]
    if not train_files:
        raise ValueError("No VTP files found in training folder")
    
    # Collect all training pressures
    # Read train.txt to get list of training VTP files
    train_txt_path = os.path.join(train_folder, '1_VTK_surface/train.txt')
    with open(train_txt_path, 'r') as f:
        folder_names = [line.strip() for line in f if line.strip()]
    
    # Construct full paths to VTP files
    train_files = []
    for folder_name in folder_names:
        vtp_file = os.path.join(train_folder, '1_VTK_surface', folder_name, f'{folder_name}.vtp')
        if os.path.exists(vtp_file):
            train_files.append(vtp_file)
        else:
            print(f"Warning: VTP file not found: {vtp_file}")
    
    if not train_files:
        raise ValueError("No VTP files found from train.txt list")
    
    train_pressures = np.concatenate([get_pressures(f) for f in train_files])
    
    # Compute bin edges and normalized histogram (PDF)
    bin_edges = np.histogram_bin_edges(train_pressures, bins=num_bins)
    train_hist, _ = np.histogram(train_pressures, bins=bin_edges, density=True)
    
    # Save histogram and bin edges
    np.savez(train_save_dir + output_file, hist=train_hist, edges=bin_edges)
    print(f"Training PDF saved to {train_save_dir + output_file}")

# Part 2: Compute JS divergence confidence score for a test VTP file
def compute_js_score(test_file, train_dist_file='train_dist.npz', pressure_field='pressure'):
    # Load training distribution
    data = np.load(train_save_dir+train_dist_file)
    train_hist = data['hist']
    bin_edges = data['edges']
    
    # Get test pressures
    test_pressures = get_pressures(test_file, pressure_field)
    
    # Compute test histogram with same bins
    test_hist, _ = np.histogram(test_pressures, bins=bin_edges, density=True)
    
    # Add small epsilon for smoothing to avoid log(0)
    epsilon = 1e-10
    train_hist = train_hist + epsilon
    test_hist = test_hist + epsilon
    train_hist /= train_hist.sum()
    test_hist /= test_hist.sum()
    
    # Compute JS divergence
    div = jensenshannon(train_hist, test_hist)
    
    # Confidence score (higher = better alignment, e.g., 1 - div)
    score = 1 - div
    return score

# Part 2b: Compute JS scores for both GT and predicted pressures
def compute_js_score_both(test_pressures, train_dist_file='train_dist.npz'):
    # Load training distribution
    data = np.load(train_save_dir+train_dist_file)
    train_hist = data['hist']
    bin_edges = data['edges']
    
    # Compute test histogram with same bins
    test_hist, _ = np.histogram(test_pressures, bins=bin_edges, density=True)
    
    # Add small epsilon for smoothing to avoid log(0)
    epsilon = 1e-10
    train_hist = train_hist + epsilon
    test_hist = test_hist + epsilon
    train_hist /= train_hist.sum()
    test_hist /= test_hist.sum()
    
    # Compute JS divergence
    div = jensenshannon(train_hist, test_hist)
    
    # Confidence score (higher = better alignment, e.g., 1 - div)
    score = 1 - div
    return score

# Part 3: Compute JS scores for test cases from test.txt, using VTK files in test directory
# Computes scores for both GT and predicted pressure
def compute_test_scores_from_vtk(vtk_test_dir, train_dist_file='train_dist.npz', output_file='test_js_scores.txt'):
    # Read test.txt to get list of test case names
    test_txt_path = os.path.join(train_folder, '1_VTK_surface/test.txt')
    with open(test_txt_path, 'r') as f:
        test_case_names = [line.strip() for line in f if line.strip()]
    
    if not test_case_names:
        raise ValueError("No test cases found in test.txt")
    
    # Save output file one directory behind the test directory
    output_path = os.path.join(os.path.dirname(test_dir+"/../"), output_file)
    
    # Compute scores for each test case
    gt_scores = []
    pred_scores = []
    
    with open(output_path, 'w') as f:
        f.write("Test_File\tGT_JS_Score\tPred_JS_Score\n")
        for test_case in test_case_names:
            # Look for corresponding VTK file in the test directory
            vtk_file = f"{test_case}.vtk"
            vtk_path = os.path.join(vtk_test_dir, vtk_file)
            
            if os.path.exists(vtk_path):
                try:
                    # Get both GT and predicted pressures
                    gt_pressure, pred_pressure = get_both_pressures(vtk_path)
                    
                    gt_score = None
                    pred_score = None
                    
                    # Compute GT score if available
                    if gt_pressure is not None:
                        gt_score = compute_js_score_both(gt_pressure, train_dist_file)
                        gt_scores.append(gt_score)
                        print(f"GT Score for {test_case}: {gt_score:.6f}")
                    
                    # Compute Pred score if available
                    if pred_pressure is not None:
                        pred_score = compute_js_score_both(pred_pressure, train_dist_file)
                        pred_scores.append(pred_score)
                        print(f"Pred Score for {test_case}: {pred_score:.6f}")
                    
                    # Write to file
                    gt_str = f"{gt_score:.6f}" if gt_score is not None else "N/A"
                    pred_str = f"{pred_score:.6f}" if pred_score is not None else "N/A"
                    f.write(f"{test_case}\t{gt_str}\t{pred_str}\n")
                    
                except Exception as e:
                    print(f"Error computing score for {test_case}: {e}")
                    f.write(f"{test_case}\tERROR: {e}\tERROR: {e}\n")
            else:
                print(f"Warning: VTK file not found for test case {test_case}: {vtk_path}")
                f.write(f"{test_case}\tERROR: VTK file not found\tERROR: VTK file not found\n")
    
    print(f"\nTest scores saved to {output_path}")
    
    # Print statistics for GT scores
    if gt_scores:
        print(f"\nGT Pressure Statistics:")
        print(f"Average GT JS Score: {np.mean(gt_scores):.6f}")
        print(f"Std GT JS Score: {np.std(gt_scores):.6f}")
        print(f"Min GT JS Score: {np.min(gt_scores):.6f}")
        print(f"Max GT JS Score: {np.max(gt_scores):.6f}")
    
    # Print statistics for Pred scores
    if pred_scores:
        print(f"\nPredicted Pressure Statistics:")
        print(f"Average Pred JS Score: {np.mean(pred_scores):.6f}")
        print(f"Std Pred JS Score: {np.std(pred_scores):.6f}")
        print(f"Min Pred JS Score: {np.min(pred_scores):.6f}")
        print(f"Max Pred JS Score: {np.max(pred_scores):.6f}")
    
    return gt_scores, pred_scores

# Part 4: Original function for VTP files from test.txt (kept for compatibility)
def compute_test_scores(test_folder, train_dist_file='train_dist.npz', output_file='test_js_scores.txt'):
    # Read test.txt to get list of test VTP files
    test_txt_path = os.path.join(test_folder, '1_VTK_surface/test.txt')
    with open(test_txt_path, 'r') as f:
        folder_names = [line.strip() for line in f if line.strip()]
    
    # Construct full paths to VTP files
    test_files = []
    for folder_name in folder_names:
        vtp_file = os.path.join(test_folder, '1_VTK_surface', folder_name, f'{folder_name}.vtp')
        if os.path.exists(vtp_file):
            test_files.append((folder_name, vtp_file))
        else:
            print(f"Warning: VTP file not found: {vtp_file}")
    
    if not test_files:
        raise ValueError("No VTP files found from test.txt list")
    
    # Compute scores for each test file
    scores = []
    with open(test_dir+"/../test_js_scores.txt", 'w') as f:
        f.write("Test_File\tJS_Score\n")
        for folder_name, vtp_file in test_files:
            try:
                score = compute_js_score(vtp_file, train_dist_file)
                scores.append(score)
                f.write(f"{folder_name}\t{score:.6f}\n")
                print(f"Score for {folder_name}: {score:.6f}")
            except Exception as e:
                print(f"Error computing score for {folder_name}: {e}")
                f.write(f"{folder_name}\tERROR: {e}\n")

    print(f"\nTest scores saved to {test_dir}/../test_js_scores.txt")
    print(f"Average JS Score: {np.mean(scores):.6f}")
    print(f"Std JS Score: {np.std(scores):.6f}")
    print(f"Min JS Score: {np.min(scores):.6f}")
    print(f"Max JS Score: {np.max(scores):.6f}")
    
    return scores

# Example Usage:
# Step 1: Compute training PDF
#compute_training_pdf(train_folder)

# Step 2: Compute scores for VTK files in test directory
test_scores = compute_test_scores_from_vtk(test_dir)