#!/usr/bin/env python3
"""
Standalone script to visualize saved denoising results.

This script allows you to generate results on a server and then visualize them locally.

Usage:
    python visualize_saved_results.py --saved_data path/to/denoising_results.pt --output path/to/output.png

Example:
    python visualize_saved_results.py --saved_data results/denoising_results.pt --output visualization.png --method umap
"""

import argparse
import os
import sys
from pathlib import Path

from mixed_diffusion.visualize import visualize_denoising


def main():
    parser = argparse.ArgumentParser(
        description="Visualize saved denoising results",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    # Basic visualization with separate subplots
    python visualize_saved_results.py --saved_data results/denoising_results.pt
    
    # Use UMAP instead of PCA and show plots
    python visualize_saved_results.py --saved_data results/denoising_results.pt --visualization_method umap --visualize
    
    # Joint visualization (all datasets in same space) 
    python visualize_saved_results.py --saved_data results/denoising_results.pt --joint_visualization
    
    # Interactive web visualization with browser opening
    python visualize_saved_results.py --saved_data results/denoising_results.pt --visualize --open_browser
    
    # Customize PCA components for UMAP preprocessing
    python visualize_saved_results.py --saved_data results/denoising_results.pt --visualization_method umap --pca_components 25
        """,
    )

    parser.add_argument(
        "--saved_data",
        type=str,
        required=True,
        help="Path to the folder containing the denoising results file (.pt)",
    )

    parser.add_argument(
        "--visualization_method",
        type=str,
        choices=["pca", "umap", "python_umap"],
        default="pca",
        help="Dimensionality reduction method for visualization (default: pca)",
    )

    parser.add_argument(
        "--visualize", action="store_true", help="Display the plot in browser/viewer"
    )

    parser.add_argument(
        "--joint_visualization",
        action="store_true",
        help="Show all datasets in joint visualization space",
    )

    parser.add_argument(
        "--pca_components",
        type=int,
        default=50,
        help="Number of PCA components before UMAP (default: 50)",
    )

    parser.add_argument(
        "--open_browser",
        action="store_true",
        help="Automatically open the HTML visualization in the browser",
    )
    parser.add_argument(
        "--test_noise_level",
        type=float,
        default=0.1,
        help="Noise level to use for testing (default: 0.1)",
    )
    parser.add_argument(
        "--result_dir",
        type=str,
        default=None,
        help="Directory to save results (default: same as saved_data directory)",
    )
    parser.add_argument(
        "--consolidate_rare_cells",
        action="store_true",
        help="Consolidate rare cell types into 'other' category (useful for CITEseq data)",
    )
    parser.add_argument(
        "--fit_umap_on_train",
        action="store_true",
        help="Fit UMAP on training data if available (requires x_train in saved data)",
    )

    args = parser.parse_args()

    output_path = args.saved_data.rstrip("/") + "/viz.png"

    args.saved_data = args.saved_data.rstrip("/") + "/denoising_results.pt"

    # Set result_dir if not provided
    if args.result_dir is None:
        args.result_dir = os.path.dirname(args.saved_data)

    # Check if saved data file exists
    if not os.path.exists(args.saved_data):
        print(f"Error: Saved data file not found: {args.saved_data}")
        sys.exit(1)

    print(f"Loading saved data from: {args.saved_data}")
    print(f"Using visualization method: {args.visualization_method}")
    print(f"Output will be saved to: {output_path}")
    if args.joint_visualization:
        print("Using joint visualization (all datasets in same space)")
    else:
        print("Using separate subplots for each dataset")

    try:
        visualize_denoising(
            saved_results=args.saved_data,
            output_path=output_path,
            args=args,
        )

        print("✅ Visualization completed successfully!")

        # Open browser if requested
        if args.open_browser:
            html_path = output_path.replace(".png", ".html")
            if os.path.exists(html_path):
                import webbrowser

                print(f"🌐 Opening {html_path} in browser...")
                webbrowser.open(f"file://{os.path.abspath(html_path)}")
            else:
                print(f"⚠️  HTML file not found: {html_path}")

    except Exception as e:
        print(f"❌ Error creating visualization: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()
