src/synthetic_generation/continuous_generation.py

import argparse
import logging
import os
import random
import signal
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.feather as feather

from src.synthetic_generation.anomalies.anomaly_generator_wrapper import (
    AnomalyGeneratorWrapper,
)
from src.synthetic_generation.audio_generators.financial_volatility_wrapper import (
    FinancialVolatilityAudioWrapper,
)
from src.synthetic_generation.audio_generators.multi_scale_fractal_wrapper import (
    MultiScaleFractalAudioWrapper,
)
from src.synthetic_generation.audio_generators.network_topology_wrapper import (
    NetworkTopologyAudioWrapper,
)
from src.synthetic_generation.audio_generators.stochastic_rhythm_wrapper import (
    StochasticRhythmAudioWrapper,
)
from src.synthetic_generation.cauker.cauker_generator_wrapper import (
    CauKerGeneratorWrapper,
)
from src.synthetic_generation.forecast_pfn_prior.forecast_pfn_generator_wrapper import (
    ForecastPFNGeneratorWrapper,
)
from src.synthetic_generation.generator_params import (
    AnomalyGeneratorParams,
    CauKerGeneratorParams,
    FinancialVolatilityAudioParams,
    ForecastPFNGeneratorParams,
    GPGeneratorParams,
    KernelGeneratorParams,
    MultiScaleFractalAudioParams,
    NetworkTopologyAudioParams,
    OrnsteinUhlenbeckProcessGeneratorParams,
    SawToothGeneratorParams,
    SineWaveGeneratorParams,
    SpikesGeneratorParams,
    StepGeneratorParams,
    StochasticRhythmAudioParams,
)
from src.synthetic_generation.gp_prior.gp_generator_wrapper import GPGeneratorWrapper
from src.synthetic_generation.kernel_synth.kernel_generator_wrapper import (
    KernelGeneratorWrapper,
)
from src.synthetic_generation.ornstein_uhlenbeck_process.ou_generator_wrapper import (
    OrnsteinUhlenbeckProcessGeneratorWrapper,
)
from src.synthetic_generation.sawtooth.sawtooth_generator_wrapper import (
    SawToothGeneratorWrapper,
)
from src.synthetic_generation.sine_waves.sine_wave_generator_wrapper import (
    SineWaveGeneratorWrapper,
)
from src.synthetic_generation.spikes.spikes_generator_wrapper import (
    SpikesGeneratorWrapper,
)
from src.synthetic_generation.steps.step_generator_wrapper import StepGeneratorWrapper


class TimeSeriesDatasetManager:
    """Manages writing time series data to disk in batches, safe for parallel runs."""

    def __init__(self, output_path: str, batch_size: int = 2**16):
        self.output_path = Path(output_path)
        self.output_path.mkdir(parents=True, exist_ok=True)
        self.batches_dir = self.output_path
        self.batch_size = batch_size
        self.series_counter = 0

        self.schema = pa.schema(
            [
                ("series_id", pa.int64()),
                ("values", pa.list_(pa.list_(pa.float64()))),
                ("length", pa.int32()),
                ("num_channels", pa.int32()),
                ("generator_type", pa.string()),
                ("start", pa.timestamp("ns")),
                ("frequency", pa.string()),
                ("generation_timestamp", pa.timestamp("ns")),
            ]
        )
        self._initialize_state()

    def _initialize_state(self) -> None:
        """Initializes state by scanning existing files to count total series."""
        existing_batches = sorted(self.batches_dir.glob("batch_*.arrow"))
        total_series = 0
        if not existing_batches:
            logging.info("No existing batches found. Starting from scratch.")
        else:
            for batch_file in existing_batches:
                try:
                    batch_table = feather.read_table(batch_file)
                    total_series += len(batch_table)
                except Exception as e:
                    logging.warning(f"Error reading {batch_file}: {e}, skipping.")
        self.series_counter = total_series
        logging.info(f"Found {self.series_counter} existing series in dataset.")

    def get_current_series_count(self) -> int:
        """Returns the total number of series found on disk at initialization."""
        return self.series_counter

    def append_batch(self, batch_data: list[dict[str, Any]]) -> None:
        """Appends a batch to a new file using an atomic rename for parallel safety."""
        if not batch_data:
            return

        try:
            arrays = []
            for field in self.schema:
                field_name = field.name
                if field_name in ["start", "generation_timestamp"]:
                    timestamps = [d[field_name] for d in batch_data]
                    arrays.append(pa.array([t.value for t in timestamps], type=pa.timestamp("ns")))
                else:
                    arrays.append(pa.array([d[field_name] for d in batch_data]))
            new_table = pa.Table.from_arrays(arrays, schema=self.schema)
        except Exception as e:
            logging.error(f"Error creating Arrow table: {e}")
            raise

        tmp_path = None
        try:
            with tempfile.NamedTemporaryFile(delete=False, dir=self.batches_dir, suffix=".arrow.tmp") as tmp:
                tmp_path = tmp.name
                feather.write_feather(new_table, tmp_path)

            max_retries = 20
            for _ in range(max_retries):
                existing = self.batches_dir.glob("batch_*.arrow")
                batch_nums = [int(p.stem.split("_")[1]) for p in existing if p.stem.split("_")[1].isdigit()]
                next_num = max(batch_nums) + 1 if batch_nums else 0
                target_path = self.batches_dir / f"batch_{next_num:08d}.arrow"
                try:
                    os.rename(tmp_path, target_path)
                    self.series_counter += len(batch_data)
                    logging.info(f"Saved {target_path.name} with {len(batch_data)} series.")
                    return
                except FileExistsError:
                    logging.warning(f"Race condition on {target_path.name}. Retrying...")
                    time.sleep(random.uniform(0.1, 1.0))

            raise OSError("Failed to write batch due to file conflicts.")
        finally:
            if tmp_path and os.path.exists(tmp_path):
                os.remove(tmp_path)


class GeneratorWrapper:
    def __init__(
        self,
        generator_type: str,
        length: int = 2048,
        global_seed: int = 42,
        num_channels: int | None = None,
    ):
        self.generator_type = generator_type
        self.length = length
        self.is_multivariate = generator_type.lower() in [
            "cauker_multivariate",
        ]
        self.explode_multivariate_to_univariate = generator_type.lower() == "cauker_univariate"
        self._explode_channels = 0

        # Create appropriate parameter object and wrapper
        if generator_type.lower() == "gp":
            params = GPGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = GPGeneratorWrapper(params)
        elif generator_type.lower() == "kernel":
            params = KernelGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = KernelGeneratorWrapper(params)
        elif generator_type.lower() == "forecast_pfn":
            params = ForecastPFNGeneratorParams(
                global_seed=global_seed,
                length=length,
                max_absolute_spread=500.0,
                max_absolute_value=500.0,
            )
            self.wrapper = ForecastPFNGeneratorWrapper(params)
        elif generator_type.lower() == "sinewave":
            params = SineWaveGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = SineWaveGeneratorWrapper(params)
        elif generator_type.lower() == "sawtooth":
            params = SawToothGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = SawToothGeneratorWrapper(params)
        elif generator_type.lower() == "cauker_univariate":
            params = CauKerGeneratorParams(
                global_seed=global_seed,
                length=length,
                num_channels=6,
            )
            self.wrapper = CauKerGeneratorWrapper(params)
            self._explode_channels = 6
        elif generator_type.lower() == "cauker_multivariate":
            effective_channels = (
                int(num_channels) if num_channels is not None else CauKerGeneratorParams().num_channels  # type: ignore[arg-type]
            )
            params = CauKerGeneratorParams(
                global_seed=global_seed,
                length=length,
                num_channels=effective_channels,
                num_nodes=effective_channels,
            )
            self.wrapper = CauKerGeneratorWrapper(params)
        elif generator_type.lower() == "step":
            params = StepGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = StepGeneratorWrapper(params)
        elif generator_type.lower() == "spike":
            params = SpikesGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = SpikesGeneratorWrapper(params)
        elif generator_type.lower() == "anomaly":
            params = AnomalyGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = AnomalyGeneratorWrapper(params)
        elif generator_type.lower() == "ou_process":
            params = OrnsteinUhlenbeckProcessGeneratorParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = OrnsteinUhlenbeckProcessGeneratorWrapper(params)
        elif generator_type.lower() == "audio_financial_volatility":
            params = FinancialVolatilityAudioParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = FinancialVolatilityAudioWrapper(params)
        elif generator_type.lower() == "audio_multi_scale_fractal":
            params = MultiScaleFractalAudioParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = MultiScaleFractalAudioWrapper(params)
        elif generator_type.lower() == "audio_stochastic_rhythm":
            params = StochasticRhythmAudioParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = StochasticRhythmAudioWrapper(params)
        elif generator_type.lower() == "audio_network_topology":
            params = NetworkTopologyAudioParams(
                global_seed=global_seed,
                length=length,
            )
            self.wrapper = NetworkTopologyAudioWrapper(params)
        else:
            raise ValueError(f"Unsupported generator type: {generator_type}")

    def generate_batch(self, batch_size: int, start_seed: int) -> list[dict[str, Any]]:
        """Generate a batch of time series using the wrapper's batch generation."""
        try:
            if self.explode_multivariate_to_univariate and self._explode_channels > 0:
                base_batch_size = int(np.ceil(batch_size / self._explode_channels))
                container = self.wrapper.generate_batch(batch_size=base_batch_size, seed=start_seed)
            else:
                container = self.wrapper.generate_batch(batch_size=batch_size, seed=start_seed)

            batch_data = []
            container_batch_size = container.values.shape[0]
            for i in range(container_batch_size):
                series_id_base = start_seed + i

                if self.explode_multivariate_to_univariate:
                    series_data = container.values[i]
                    if series_data.ndim != 2:
                        raise ValueError("Expected multivariate data for CauKer univariate mode")
                    num_channels = series_data.shape[1]
                    for channel in range(num_channels):
                        channel_values = self._ensure_proper_format(series_data[:, channel])
                        values_list = [channel_values.tolist()]
                        batch_data.append(
                            {
                                "series_id": series_id_base * 1_000 + channel,
                                "values": values_list,
                                "length": len(channel_values),
                                "num_channels": 1,
                                "generator_type": self.generator_type,
                                "start": pd.Timestamp(container.start[i]),
                                "frequency": container.frequency[i].value,
                                "generation_timestamp": pd.Timestamp.now(),
                            }
                        )
                    continue
                elif self.is_multivariate:
                    series_data = container.values[i]
                    num_channels = series_data.shape[1]
                    values_list = [self._ensure_proper_format(series_data[:, c]).tolist() for c in range(num_channels)]
                    seq_length = len(values_list[0])
                else:
                    values = self._ensure_proper_format(container.values[i, :])
                    values_list = [values.tolist()]
                    num_channels = 1
                    seq_length = len(values)

                batch_data.append(
                    {
                        "series_id": series_id_base,
                        "values": values_list,
                        "length": seq_length,
                        "num_channels": num_channels,
                        "generator_type": self.generator_type,
                        "start": pd.Timestamp(container.start[i]),
                        "frequency": container.frequency[i].value,
                        "generation_timestamp": pd.Timestamp.now(),
                    }
                )

            if self.explode_multivariate_to_univariate:
                batch_data = batch_data[:batch_size]

            return batch_data

        except Exception as e:
            logging.error(f"Error generating batch: {e}")
            return []

    def _ensure_proper_format(self, values: Any) -> np.ndarray:
        values = np.asarray(values).flatten()
        if len(values) != self.length:
            logging.warning(f"Generated series length {len(values)} != expected {self.length}. Padding/truncating.")
            if len(values) > self.length:
                values = values[: self.length]
            else:
                values = np.pad(values, (0, self.length - len(values)), mode="constant")
        return values.astype(np.float64)


class ContinuousGenerator:
    def __init__(
        self,
        generator_wrapper: GeneratorWrapper,
        dataset_manager: TimeSeriesDatasetManager,
        batch_size: int = 2**16,
        run_id: int = 0,
    ):
        self.generator_wrapper = generator_wrapper
        self.dataset_manager = dataset_manager
        self.batch_size = batch_size
        self.run_id = run_id
        self.series_in_run = 0
        self.partial_batch_data: list[dict[str, Any]] = []
        self.shutting_down = False
        logging.info(f"Generator initialized for run_id: {self.run_id}")

    def _setup_signal_handlers(self) -> None:
        """Sets up signal handlers for graceful shutdown."""
        self.original_sigint = signal.getsignal(signal.SIGINT)
        self.original_sigterm = signal.getsignal(signal.SIGTERM)

        def graceful_shutdown(signum, frame):
            if self.shutting_down:
                return
            self.shutting_down = True
            logging.warning(f"\nSignal {signal.Signals(signum).name} received. Shutting down.")
            if self.partial_batch_data:
                logging.info(f"Saving incomplete batch of {len(self.partial_batch_data)} series...")
                try:
                    self.dataset_manager.append_batch(self.partial_batch_data)
                except Exception as e:
                    logging.error(f"Could not save partial batch on exit: {e}")
            sys.exit(0)

        signal.signal(signal.SIGINT, graceful_shutdown)
        signal.signal(signal.SIGTERM, graceful_shutdown)

    def run_continuous(self, num_batches_to_generate: int) -> None:
        """Runs the generation loop, creating chunks and saving batches."""
        self._setup_signal_handlers()
        logging.info(f"Job starting. Goal: {num_batches_to_generate} new batches.")
        start_time = time.time()
        batches_completed = 0

        while batches_completed < num_batches_to_generate:
            if self.shutting_down:
                logging.info("Shutdown signal caught, stopping generation.")
                break

            chunk_size = min(64, self.batch_size - len(self.partial_batch_data))

            # Create a seed that fits in uint32 range by combining run_id and series count
            # Use modulo to ensure it stays within valid range
            series_id_start = (self.run_id + self.series_in_run) % (2**32)

            new_chunk = self.generator_wrapper.generate_batch(batch_size=chunk_size, start_seed=series_id_start)

            if not new_chunk:
                logging.error("Generator failed to produce data. Stopping job.")
                break

            self.partial_batch_data.extend(new_chunk)
            self.series_in_run += len(new_chunk)

            if len(self.partial_batch_data) >= self.batch_size:
                batch_to_write = self.partial_batch_data[: self.batch_size]
                self.partial_batch_data = self.partial_batch_data[self.batch_size :]
                self.dataset_manager.append_batch(batch_to_write)
                batches_completed += 1

                elapsed = time.time() - start_time
                series_per_sec = (batches_completed * self.batch_size) / elapsed if elapsed > 0 else 0
                print(
                    f"✓ Completed batch {batches_completed}/{num_batches_to_generate} in job | "
                    f"Total Series in DS: {self.dataset_manager.series_counter:,} | "
                    f"Rate: {series_per_sec:.1f}/s"
                )

        if not self.shutting_down and self.partial_batch_data:
            logging.info(f"Job finished. Saving final partial batch of {len(self.partial_batch_data)}.")
            self.dataset_manager.append_batch(self.partial_batch_data)


def setup_logging(verbose: bool = False) -> None:
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[logging.StreamHandler(sys.stdout)],
    )


def main():
    parser = argparse.ArgumentParser(
        description="Continuous time series generation script",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--generator",
        type=str,
        required=True,
        choices=[
            "forecast_pfn",
            "gp",
            "kernel",
            "cauker_univariate",
            "cauker_multivariate",
            "sinewave",
            "sawtooth",
            "step",
            "spike",
            "anomaly",
            "ou_process",
            "audio_financial_volatility",
            "audio_multi_scale_fractal",
            "audio_stochastic_rhythm",
            "audio_network_topology",
        ],
        help="Type of generator to use",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="Output directory for datasets",
    )
    parser.add_argument("--length", type=int, default=2048, help="Length of each time series")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=16384,
        help="Number of series per batch file",
    )
    parser.add_argument(
        "--num-batches",
        type=int,
        default=100,
        help="Number of batches to generate in this job run",
    )
    parser.add_argument(
        "--num-channels",
        type=int,
        help="Number of channels for multivariate generators (cauker_multivariate)",
    )
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
    args = parser.parse_args()

    setup_logging(args.verbose)

    # Use a high-precision timestamp for a unique run ID and a compatible seed
    run_id = time.time_ns()
    global_seed = run_id % (2**32)
    logging.info(f"Using unique run ID: {run_id} (Seed: {global_seed})")

    gen_name = args.generator.lower()
    if gen_name in ["cauker_multivariate"]:
        if args.num_channels is None or args.num_channels < 2:
            logging.error("--num-channels (>=2) is required for multivariate generators")
            sys.exit(2)
        dataset_dir_name = f"cauker_{args.num_channels}_variates"
    else:
        dataset_dir_name = args.generator

    output_path = Path(args.output_dir) / dataset_dir_name

    try:
        generator_wrapper = GeneratorWrapper(
            generator_type=args.generator,
            length=args.length,
            global_seed=global_seed,
            num_channels=args.num_channels,
        )
        dataset_manager = TimeSeriesDatasetManager(str(output_path), batch_size=args.batch_size)
        continuous_gen = ContinuousGenerator(
            generator_wrapper=generator_wrapper,
            dataset_manager=dataset_manager,
            batch_size=args.batch_size,
            run_id=run_id,
        )
        continuous_gen.run_continuous(num_batches_to_generate=args.num_batches)
        logging.info("Generation job completed successfully!")

    except Exception as e:
        logging.error(f"Fatal error: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()