from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

CURVE_KEYS = ("aif", "tac", "time")


def batch_to_dataframe(batch: dict[str, np.ndarray]) -> pd.DataFrame:
    n = batch["tac"].shape[0]
    cols = {}

    for key, value in batch.items():
        if key in CURVE_KEYS:
            continue
        arr = np.asarray(value)
        if arr.shape[0] != n:
            raise ValueError(
                f"Column '{key}' has {arr.shape[0]} rows, expected {n} rows."
            )
        cols[key] = arr

    for key in CURVE_KEYS:
        curve = np.asarray(batch[key], dtype=np.float32, order="C")
        if curve.ndim != 2 or curve.shape[0] != n:
            raise ValueError(
                f"Curve '{key}' must have shape (N, T). Got {tuple(curve.shape)}."
            )
        t = curve.shape[1]
        for j in range(t):
            cols[f"{key}_{j:03d}"] = curve[:, j]

    return pd.DataFrame(cols)


class SimulationParquetWriter:
    def __init__(
        self,
        out_dir,
        *,
        file_name: str = "simulations.parquet",
        compression: str = "zstd",
    ):
        self.out_dir = Path(out_dir)
        self.path = self.out_dir / file_name
        self.compression = compression
        self._writer = None

    def __enter__(self):
        self.out_dir.mkdir(parents=True, exist_ok=True)
        if self.path.exists():
            self.path.unlink()
        return self

    def write_batch(self, frame: pd.DataFrame):
        table = pa.Table.from_pandas(frame, preserve_index=False)
        if self._writer is None:
            self._writer = pq.ParquetWriter(
                str(self.path),
                table.schema,
                compression=self.compression,
                use_dictionary=True,
                write_statistics=True,
            )
        self._writer.write_table(table)

    def close(self):
        if self._writer is not None:
            self._writer.close()
            self._writer = None

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()


__all__ = ["SimulationParquetWriter", "batch_to_dataframe"]
