# FILE: a3d/export.py
from __future__ import annotations

import csv
from typing import Any, Dict, List, Optional

import numpy as np


def export_csv(path: str, rows: List[Dict[str, Any]]) -> None:
    if not rows:
        with open(path, "w", newline="") as f:
            f.write("")
        return
    headers = sorted(rows[0].keys())
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=headers)
        w.writeheader()
        for r in rows:
            w.writerow(r)


def export_parquet(path: str, rows: List[Dict[str, Any]]) -> None:
    try:
        import pyarrow as pa
        import pyarrow.parquet as pq
    except Exception as e:
        raise RuntimeError("pyarrow not available") from e
    if not rows:
        table = pa.Table.from_pydict({})
    else:
        keys = sorted(rows[0].keys())
        cols = {k: [r.get(k, None) for r in rows] for k in keys}
        table = pa.Table.from_pydict(cols)
    pq.write_table(table, path)


def export_fixed_point_header(
    path: str,
    arrays: Dict[str, np.ndarray],
    scales: Dict[str, float],
    meta: Optional[Dict[str, Any]] = None,
) -> None:
    """
    Write a C header with fixed-point int16 arrays.
    - arrays: dict of {name: np.ndarray}
    - scales: dict of {name: float} such that real_value ≈ int16_value * scale
    """

    def quantize(a: np.ndarray, scale: float) -> np.ndarray:
        q = np.round(a / scale).astype(np.int64)
        q = np.clip(q, -32768, 32767).astype(np.int16)
        return q

    with open(path, "w") as f:
        f.write("/* Autogenerated by Aegis export_fixed_point_header */\n")
        f.write("#pragma once\n#include <stdint.h>\n\n")
        if meta:
            for k, v in meta.items():
                f.write(f"/* {k}: {v} */\n")
        for name, arr in arrays.items():
            sc = float(scales[name])
            q = quantize(np.asarray(arr, dtype=np.float64), sc)
            shape = q.shape
            flat = q.flatten()
            f.write(f"#define {name.upper()}_SCALE {sc:.12g}\n")
            for ax, sz in enumerate(shape):
                f.write(f"#define {name.upper()}_DIM{ax} {int(sz)}\n")
            f.write(f"static const int16_t {name}[] = {{\n")
            for _i, val in enumerate(flat):
                f.write(f"  {int(val)},\n")
            f.write("};\n\n")
