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.")