#!/usr/bin/env python
"""
Start TensorBoard for monitoring neural network training
"""

import os
import sys
import subprocess
import webbrowser
from time import sleep

def start_tensorboard(logdir="NN/models/saved/logs", port=6006, open_browser=True):
    """
    Start TensorBoard in a subprocess
    
    Args:
        logdir: Directory containing TensorBoard logs
        port: Port to run TensorBoard on
        open_browser: Whether to open a browser automatically
    """
    # Make sure the log directory exists
    os.makedirs(logdir, exist_ok=True)
    
    # Create command
    cmd = [
        sys.executable,
        "-m",
        "tensorboard.main",
        f"--logdir={logdir}",
        f"--port={port}",
        "--bind_all"
    ]
    
    print(f"Starting TensorBoard with logs from {logdir} on port {port}")
    print(f"Command: {' '.join(cmd)}")
    
    # Start TensorBoard in a subprocess
    process = subprocess.Popen(
        cmd, 
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True
    )
    
    # Wait for TensorBoard to start up
    for line in process.stdout:
        print(line.strip())
        if "TensorBoard" in line and "http://" in line:
            # TensorBoard is running, extract the URL
            url = None
            for part in line.split():
                if part.startswith(("http://", "https://")):
                    url = part
                    break
            
            # Open browser if requested and URL found
            if open_browser and url:
                print(f"Opening TensorBoard in browser: {url}")
                webbrowser.open(url)
            
            break
    
    # Return the process for the caller to manage
    return process

if __name__ == "__main__":
    import argparse
    
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Start TensorBoard for NN training visualization")
    parser.add_argument("--logdir", default="NN/models/saved/logs", help="Directory containing TensorBoard logs")
    parser.add_argument("--port", type=int, default=6006, help="Port to run TensorBoard on")
    parser.add_argument("--no-browser", action="store_true", help="Don't open browser automatically")
    
    args = parser.parse_args()
    
    # Start TensorBoard
    process = start_tensorboard(args.logdir, args.port, not args.no_browser)
    
    try:
        # Keep the script running until Ctrl+C
        print("TensorBoard is running. Press Ctrl+C to stop.")
        while True:
            sleep(1)
    except KeyboardInterrupt:
        print("Stopping TensorBoard...")
        process.terminate()
        process.wait()