main.py

from typing import List

from matplotlib import pyplot as plt
from rtest.analysis.results_analysis import further_analysis
from rtest.methods.shared.results import SimulationResults
from rtest.config.input_config import InputConfig
from rtest.config.globals import write_js_config
from rtest.methods.tdhf.td_hf import tdhf
from rtest.methods.tddft.td_dft import td_dft_simulation
from rtest.plotting.plotter import plot_simulation_results
from rtest.results.results import save_results, save_figure
from rtest import setup_logging
import logging

# matplotlib.use('TkAgg')

if __name__ == "__main__":
  setup_logging(log_level="INFO")
  
  # Generate JavaScript config from Python globals (ensures visualization uses correct paths)
  write_js_config()
  
  ic = InputConfig.from_yaml(yaml_path="input.yaml", do_validation=True)
  # Setup Logging & Warnings
  setup_logging(log_level=ic.log_level)
  logger = logging.getLogger("rtest.main")
  
  logger.info(f"Starting simulation. Method: [bold cyan]{ic.simulation_method}[/]", extra={"markup": True})
  match ic.simulation_method:
    case "hf" | "dft":
      results: List[SimulationResults] = tdhf(ic=ic, method=ic.simulation_method)
    case "dft_not_von_neumann":
      results: SimulationResults = td_dft_simulation(ic)
    case _:
      raise ValueError(f"Unknown simulation method: {ic.simulation_method}. Expected 'hf', 'dft', or 'dft_not_von_neumann'.")
    
  logger.info("Simulation complete.")
  logger.info("Starting analysis if requested...")
  analysis_results = further_analysis(ic, results)
  
  logger.info("Saving results...")
  save_results(analysis_results, results, ic)
  
  logger.info("Generating plots...")
  fig: plt.Figure = plot_simulation_results(analysis_results, ic)
  
  # Only show plot if we have an interactive backend
  backend = plt.get_backend()
  if backend not in ['agg', 'png', 'svg', 'pdf', 'ps']:
    try:
      plt.show()
    except Exception as e:
      logger.debug(f"Could not show plot interactively: {e}")
  else:
    logger.info(f"Not showing plot interactively due to non-interactive matplotlib backend: {backend}")
  
  logger.info("Saving figure...")
  save_figure(fig, ic, filename="simulation_results.png", dpi=1000)

# # Access the results
# print(f"Simulation results: {results}")

# # Access specific data from the results
# print(f"Calculated induced dipole x: {results[0].data_handler.get_data('absorption_induced_dipole_x')}")