import numpy as np
import os
import minipandas as mpd   # the teaching version of pandas :)

print("=== Creating test MiniFrames ===")

# Mix of scalar + 1D + 2D columns
mf1 = mpd.MiniFrame({
    "run": 42,                             # scalar → broadcast
    "event": np.arange(5),                 # 1D
    "energy": np.random.rand(5, 3),        # 2D fixed-size
})

mf2 = mpd.MiniFrame({
    "run": 43,
    "event": np.arange(5, 10),
    "energy": np.random.rand(5, 3),
})

print(mf1)
print("\nScalar columns:", mf1.scalar_cols)
print("Vector columns:", mf1.vector_cols)
print("Vector length(energy):", mf1.vector_length("energy"))

print("\n=== Slicing and access ===")
print("First 3 rows:")
print(mf1[:3])
print("\nEnergy column sample:")
print(mf1["energy"][:2])

print("\n=== Concatenating like pandas.concat ===")
mf_all = mpd.concat([mf1, mf2])
print(mf_all)
print("Shape:", mf_all.shape)

print("\n=== Converting to pandas DataFrame ===")
df = mf_all.to_pandas()
print(df.head())

print("\n=== Saving and reloading with pickle ===")
path = "minipandas_test.pkl"
mf_all.to_pickle(path)
mf_loaded = mpd.MiniFrame.read_pickle(path)
print("Reloaded MiniFrame:")
print(mf_loaded)
os.remove(path)

print("\n=== Describe ===")
desc = mf_all.describe()
print(desc)

print("\nAll basic minipandas features appear to work correctly.")



# --- Create a simple MiniFrame ---
N, M = 3, 4
mf = mpd.MiniFrame({
    "run": 1,                              # scalar -> broadcast
    "event": np.arange(N),                 # 1D scalar column
    "energy": np.arange(N * M).reshape(N, M),  # 2D vector column
})

# --- Case 1: indiv_cols=True (default) ---
df_wide = mf.to_pandas(indiv_cols=True)
print("=== Expanded columns ===")
print(df_wide)

# Expected shape: N rows, 1 scalar column ('event'), 1 broadcasted column ('run'), and M expanded columns
assert list(df_wide.columns) == ["run", "event"] + [f"energy_{i}" for i in range(M)]
assert df_wide.shape == (N, 2 + M)
assert np.allclose(df_wide["energy_0"], [0, 4, 8])

# --- Case 2: indiv_cols=False (legacy nested format) ---
df_nested = mf.to_pandas(indiv_cols=False)
print("\n=== Nested legacy format ===")
print(df_nested)

# Expected columns: run, event, energy
assert list(df_nested.columns) == ["run", "event", "energy"]
assert isinstance(df_nested.loc[0, "energy"], list)
assert df_nested.loc[0, "energy"] == [0, 1, 2, 3]

# --- Round-trip compatibility check ---
mf2 = mpd.MiniFrame({
    "run": df_nested["run"].to_numpy(),
    "event": df_nested["event"].to_numpy(),
    "energy": np.stack(df_nested["energy"].to_numpy()),
})

assert np.allclose(mf["energy"], mf2["energy"])
print("\n✅ All checks passed successfully!")