import json
from pathlib import Path

import numpy as np
import pyarrow as pa
import pyarrow.dataset as ds

T = 42  # number of frames


def write_time_sidecar(out_dir: str, seed: int):
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    payload = {
        "seed": seed,
    }
    (out / "_metadata.json").write_text(json.dumps(payload, indent=2))


def make_table(
    sample_id: np.ndarray,
    factors: dict,  # e.g. {"protocol": ..., "noise": ...} arrays length N
    theta: dict,  # e.g. {"K1":..., "k2":...} arrays length N
    aif_params: dict,  # arrays length N
    aif: np.ndarray,  # (N, T) float32/float64
    tac: np.ndarray,  # (N, T)
    time: np.ndarray,  # (N, T)
):
    N = aif.shape[0]
    assert aif.shape == (N, T)
    assert tac.shape == (N, T)
    assert time.shape == (N, T)

    cols = {
        "sample_id": pa.array(sample_id),
    }

    # factors + params as scalar columns
    for k, v in factors.items():
        cols[k] = pa.array(v)
    for k, v in theta.items():
        cols[k] = pa.array(v)
    for k, v in aif_params.items():
        cols[k] = pa.array(v)

    # wide curve columns
    for j in range(T):
        cols[f"aif_{j:03d}"] = pa.array(aif[:, j])
        cols[f"tac_{j:03d}"] = pa.array(tac[:, j])
        cols[f"time_{j:03d}"] = pa.array(time[:, j])

    return pa.table(cols)


def write_parquet_dataset(
    table: pa.Table,
    root_dir: str,
    partition_cols=("protocol", "noise"),
    compression="zstd",
):
    root = Path(root_dir) / "parquet"
    root.mkdir(parents=True, exist_ok=True)

    ds.write_dataset(
        data=table,
        base_dir=str(root),
        format="parquet",
        partitioning=list(partition_cols),
        existing_data_behavior="overwrite_or_ignore",  # safe for repeated writes
        file_options=ds.ParquetFileFormat().make_write_options(
            compression=compression,
            use_dictionary=True,
            write_statistics=True,
        ),
        max_rows_per_file=50_000,  # tune
        max_rows_per_group=10_000,  # tune
    )


def batch_to_table(batch):
    # batch is a dict of numpy arrays length N (and curves are (N,T))
    N = batch["tac"].shape[0]
    cols = {
        "batch_id": pa.array(batch["batch_id"]),
        "sample_id": pa.array(batch["sample_id"]),
        "protocol": pa.array(batch["protocol"]),
        "noise": pa.array(batch["noise"]),
        "regime": pa.array(batch["regime"]),
        "n_frames_used": pa.array(batch["n_frames_used"]),
        # params...
        "K1": pa.array(batch["K1"]),
        "k2": pa.array(batch["k2"]),
        "k3": pa.array(batch["k3"]),
        "k4": pa.array(batch["k4"]),
        "vB": pa.array(batch["vB"]),
        # aif params...
        "Tinj_s": pa.array(batch["Tinj_s"]),
        "tau": pa.array(batch["tau"]),
        "A1": pa.array(batch["A1"]),
        "r2": pa.array(batch["r2"]),
        "r3": pa.array(batch["r3"]),
        "L1": pa.array(batch["L1"]),
        "L2": pa.array(batch["L2"]),
        "L3": pa.array(batch["L3"]),
    }

    aif = batch["aif"].astype(np.float32, copy=False)
    tac = batch["tac"].astype(np.float32, copy=False)
    time = batch["time"].astype(np.float32, copy=False)
    for j in range(T):
        cols[f"aif_{j:03d}"] = pa.array(aif[:, j])
        cols[f"tac_{j:03d}"] = pa.array(tac[:, j])
        cols[f"time_{j:03d}"] = pa.array(time[:, j])

    return pa.table(cols)


def write_batch(
    table: pa.Table,
    out_dir: str,
    partition_cols=("batch_id", "protocol", "noise", "regime"),
    compression="zstd",
):
    base_dir = str(Path(out_dir) / "parquet")
    ds.write_dataset(
        table,
        base_dir=base_dir,
        format="parquet",
        partitioning=list(partition_cols),
        partitioning_flavor="hive",
        existing_data_behavior="overwrite_or_ignore",
        file_options=ds.ParquetFileFormat().make_write_options(
            compression=compression,
            use_dictionary=True,
            write_statistics=True,
        ),
        max_rows_per_file=50_000,
        max_rows_per_group=10_000,
    )
