#!/usr/bin/env python
import os
import itertools
import glob
import sys
from Data.harmonic_oscillator.make_data2 import get_experiment_param_ranges

def generate_configs(seq_lengths, latent_sigmas, template_file, output_dir):
  """
  Generate configuration files using template substitution to preserve all original formatting.
  Uses 2-space indentation throughout.
  """
  # First, delete existing configuration files with the ho_ prefix
  existing_files = glob.glob(os.path.join(output_dir, "ho_*.yaml"))
  for file_path in existing_files:
    try:
      os.remove(file_path)
      print(f"Deleted existing file: {file_path}")
    except Exception as e:
      print(f"Error deleting {file_path}: {e}")

  # Read the template file as a string
  with open(template_file, "r") as f:
    template_content = f.read()

  # Create output directory if it doesn't exist
  os.makedirs(output_dir, exist_ok=True)

  # Generate configurations for each parameter combination
  count = 0
  for seq_length, latent_sigma in itertools.product(seq_lengths, latent_sigmas):
    # Create a copy of the template content
    content = template_content

    # Replace the specific lines with new parameter values
    content = content.replace(
      "  seq_length: 100 # Higher number means more densely sampled data",
      f"  seq_length: {seq_length} # Higher number means more densely sampled data"
    )
    content = content.replace(
      "  pred_length: 99",
      f"  pred_length: {seq_length - 1}"
    )
    content = content.replace(
      "  data_latent_sigma: 0.5 # process_noise for the data generating SDE",
      f"  data_latent_sigma: {latent_sigma} # process_noise for the data generating SDE"
    )

    # Create filename with the specified format
    latent_sigma_str = f"{latent_sigma:.2f}".replace('.', 'p')

    filename = f"ho_{seq_length}_{latent_sigma_str}.yaml"
    output_path = os.path.join(output_dir, filename)

    # Save the new configuration
    with open(output_path, "w") as f:
      f.write(content)

    print(f"Generated {output_path}")
    count += 1

  print(f"Created {count} configuration files with template substitution.")

if __name__ == "__main__":

  """
  Don't forget to run the data creation script before running this!
  python -m Data.harmonic_oscillator.make_data2
  """

  # Get parameter ranges from the experiment setup function
  param_ranges = get_experiment_param_ranges()

  # Define the parameter combinations
  seq_lengths = param_ranges['seq_length_range'].tolist()
  latent_sigmas = param_ranges['process_noise_range'].tolist()  # process noise
  # noise_stds = param_ranges['observation_noise_range'].tolist()  # observation noise

  # Template file path
  template_file = "Config/harmonic_oscillator/harmonic_oscillator.yaml"

  # Output directory where all generated files will be saved
  output_dir = "Config/harmonic_oscillator/"

  # Generate the configuration files
  generate_configs(seq_lengths, latent_sigmas, template_file, output_dir)