import ktc
from ktc import ModelConfig
import os
import numpy as np

# -------- CONFIGURATION --------
ONNX_MODEL = "/workspace/output/dynamic_liagcn_best_1_optimized_25.onnx"
X_FILE = "/workspace/calib_data/X_norm.npy"
Y_FILE = "/workspace/calib_data/y_norm.npy"
OUTPUT_DIR = "/workspace/output/dynamic_liagcn_best_1_working_optimized_25_bie"
PLATFORM = "730"

# Quantization settings
SAMPLES_PER_CLASS = 100  # Number of samples to take from each unique class
QUANT_MODE = "mix balance" 

os.makedirs(OUTPUT_DIR, exist_ok=True)

# -------- STEP 1: Balanced Data Selection --------
print(f"[Step 1] Loading data and selecting balanced samples...")

X_raw = np.load(X_FILE) # Shape: (2462, 3, 60, 17, 1)
y_raw = np.load(Y_FILE) # Shape: (2462,)

# Preprocess X immediately to match model input
# Squeeze last dim and convert to float32 as per test_stgcn_onnx.py
X_processed = np.squeeze(X_raw, axis=-1).astype(np.float32)
# X_processed = np.expand_dims(X_processed, axis=1)
# X_processed = np.repeat(X_processed, 10, axis=1)
# X_processed = np.tile(X_processed, (1, 10, 1, 1))

# Identify unique classes
unique_classes = np.unique(y_raw)
print(f"  Found {len(unique_classes)} unique classes: {unique_classes}")

selected_indices = []

for cls in unique_classes:
    # Get all indices belonging to this class
    indices = np.where(y_raw == cls)[0]
    
    # Take a subset (up to SAMPLES_PER_CLASS)
    count = min(len(indices), SAMPLES_PER_CLASS)
    selected_indices.extend(indices[:count])
    print(f"  Class {cls}: selected {count} samples")

# Create the final calibration list
# NOTE: Ensure "input" matches your ONNX input name from test_stgcn_onnx.py
INPUT_NAME = "input" 

input_mapping = {
    INPUT_NAME: [X_processed[i:i+1] for i in selected_indices]
}

print(f"\nTotal calibration samples: {len(selected_indices)}")

# -------- STEP 2: Model Configuration --------
print(f"\n[Step 2] Configuring model for {PLATFORM}...")

km = ModelConfig(
    32768,          # DEVICE_ID for KDP730
    "8b28",         # OS_ID
    PLATFORM,       # SUB_TYPE
    onnx_path=ONNX_MODEL,
    debug=True
)

# -------- STEP 3: Run Quantization --------
print(f"\n[Step 3] Running {QUANT_MODE} quantization...")

bie_path = km.analysis(
    input_mapping,
    output_dir=OUTPUT_DIR,
    threads=4,
    datapath_bitwidth_mode=QUANT_MODE
)

print(f"\n✓ Balanced BIE generated: {bie_path}")
# import ktc
# from ktc import ModelConfig
# import os
# import numpy as np

# # -------- CONFIGURATION --------
# # Ensure this matches the file name from your export script
# ONNX_MODEL = "/workspace/output/stgcn_optimized.onnx"
# OUTPUT_DIR = "/workspace/output/stgcn_optimized_bie"
# PLATFORM = "730"

# # Calibration settings
# NUM_CALIB_SAMPLES = 100 
# QUANT_MODE = "mix balance" 

# # Model Input Metadata
# INPUT_NAME = "input"
# CHANNELS = 3
# FRAMES = 64
# KEYPOINTS = 16

# os.makedirs(OUTPUT_DIR, exist_ok=True)

# # -------- STEP 1: Random Data Generation --------
# print(f"[Step 1] Generating {NUM_CALIB_SAMPLES} random samples for calibration...")

# # Create random data mimicking the expected input shape: (1, 3, 64, 16)
# # We use uniform distribution between 0 and 1 to mimic normalized coordinates
# random_data = [
#     np.random.uniform(0, 1, (1, CHANNELS, FRAMES, KEYPOINTS)).astype(np.float32)
#     for _ in range(NUM_CALIB_SAMPLES)
# ]

# input_mapping = {
#     INPUT_NAME: random_data
# }

# print(f"  Input name: {INPUT_NAME}")
# print(f"  Sample shape: {random_data[0].shape}")

# # -------- STEP 2: Model Configuration --------
# print(f"\n[Step 2] Configuring model for KDP{PLATFORM}...")

# # Device ID 32771 is standard for KL730
# km = ModelConfig(
#     32771,          
#     "8b28",         
#     PLATFORM,       
#     onnx_path=ONNX_MODEL,
#     debug=True
# )

# # -------- STEP 3: Run Quantization & Analysis --------
# print(f"\n[Step 3] Running {QUANT_MODE} quantization...")

# # This generates the .bie file which is required for the final .nef compiler
# bie_path = km.analysis(
#     input_mapping,
#     output_dir=OUTPUT_DIR,
#     threads=8, # Increased threads for faster processing if available
#     datapath_bitwidth_mode=QUANT_MODE
# )

# print(f"\n✓ BIE generated successfully: {bie_path}")
# print("Next step: Run the Kneron Compiler on this .bie file to get your .nef")